fix bug #10 for adding entities

Author: vincentlauvlwj

ktorm-core/src/main/kotlin/me/liuwj/ktorm/entity/EntityDml.kt

@@ -10,6 +10,8 @@ import me.liuwj.ktorm.schema.*
  */
 @Suppress("UNCHECKED_CAST")
 fun <E : Entity<E>> Table<E>.add(entity: E): Int {
+    entity.implementation.checkUnexpectedDiscarding(this)
+
     val assignments = findInsertColumns(entity).takeIf { it.isNotEmpty() } ?: return 0
 
     val expression = AliasRemover.visit(
@@ -66,6 +68,8 @@ private fun Table<*>.findInsertColumns(entity: Entity<*>): Map<Column<*>, Any?>
 @Suppress("UNCHECKED_CAST")
 internal fun EntityImplementation.doFlushChanges(): Int {
     val fromTable = this.fromTable?.takeIf { this.parent == null } ?: error("The entity is not associated with any table yet.")
+    checkUnexpectedDiscarding(fromTable)
+
     val primaryKey = fromTable.primaryKey ?: error("Table ${fromTable.tableName} doesn't have a primary key.")
     val assignments = findChangedColumns(fromTable).takeIf { it.isNotEmpty() } ?: return 0
 
@@ -111,26 +115,17 @@ private fun EntityImplementation.findChangedColumns(fromTable: Table<*>): Map<Co
                 var anyChanged = false
                 var curr: Any? = this
 
-                for ((i, prop) in binding.properties.withIndex()) {
+                for (prop in binding.properties) {
                     if (curr is Entity<*>) {
                         curr = curr.implementation
                     }
 
                     check(curr is EntityImplementation?)
 
-                    val changed = if (curr == null) false else prop.name in curr.changedProperties
-
-                    // Add check to avoid bug #10
-                    if (changed && i > 0) {
-                        check(curr != null)
-
-                        if (curr.fromTable != null && curr.getRoot() != this) {
-                            val propPath = binding.properties.subList(0, i + 1).joinToString(separator = ".", prefix = "this.") { it.name }
-                            throw IllegalStateException("$propPath may be unexpectedly discarded after flushChanges, please save it to database first.")
-                        }
+                    if (curr != null && prop.name in curr.changedProperties) {
+                        anyChanged = true
                     }
 
-                    anyChanged = anyChanged || changed
                     curr = curr?.getProperty(prop.name)
                 }
 
@@ -144,15 +139,6 @@ private fun EntityImplementation.findChangedColumns(fromTable: Table<*>): Map<Co
     return assignments
 }
 
-private tailrec fun EntityImplementation.getRoot(): EntityImplementation {
-    val parent = this.parent
-    if (parent == null) {
-        return this
-    } else {
-        return parent.getRoot()
-    }
-}
-
 internal fun EntityImplementation.doDiscardChanges() {
     val fromTable = this.fromTable?.takeIf { this.parent == null } ?: error("The entity is not associated with any table yet.")
 
@@ -183,6 +169,54 @@ internal fun EntityImplementation.doDiscardChanges() {
     }
 }
 
+// Add check to avoid bug #10
+private fun EntityImplementation.checkUnexpectedDiscarding(fromTable: Table<*>) {
+    for (column in fromTable.columns) {
+        val binding = column.binding?.takeIf { column is SimpleColumn } ?: continue
+
+        if (binding is NestedBinding) {
+            var curr: Any? = this
+
+            for ((i, prop) in binding.properties.withIndex()) {
+                if (curr == null) {
+                    break
+                }
+                if (curr is Entity<*>) {
+                    curr = curr.implementation
+                }
+
+                check(curr is EntityImplementation)
+
+                if (i > 0 && prop.name in curr.changedProperties && curr.fromTable != null && curr.getRoot() != this) {
+                    val propPath = binding.properties.subList(0, i + 1).joinToString(separator = ".", prefix = "this.") { it.name }
+                    throw IllegalStateException("$propPath may be unexpectedly discarded, please save it to database first.")
+                }
+
+                curr = curr.getProperty(prop.name)
+            }
+        }
+    }
+}
+
+private tailrec fun EntityImplementation.getRoot(): EntityImplementation {
+    val parent = this.parent
+    if (parent == null) {
+        return this
+    } else {
+        return parent.getRoot()
+    }
+}
+
+internal fun Entity<*>.clearChangesRecursively() {
+    implementation.changedProperties.clear()
+
+    for ((_, value) in properties) {
+        if (value is Entity<*>) {
+            value.clearChangesRecursively()
+        }
+    }
+}
+
 @Suppress("UNCHECKED_CAST")
 internal fun EntityImplementation.doDelete(): Int {
     val fromTable = this.fromTable?.takeIf { this.parent == null } ?: error("The entity is not associated with any table yet.")

ktorm-core/src/main/kotlin/me/liuwj/ktorm/entity/EntityFinding.kt

@@ -114,7 +114,7 @@ private infix fun ColumnDeclaring<*>.eq(column: ColumnDeclaring<*>): BinaryExpre
 @Suppress("UNCHECKED_CAST")
 fun <E : Entity<E>> Table<E>.createEntity(row: QueryRowSet): E {
     val entity = doCreateEntity(row, skipReferences = false) as E
-    return entity.apply { discardChanges() }
+    return entity.apply { clearChangesRecursively() }
 }
 
 /**
@@ -123,7 +123,7 @@ fun <E : Entity<E>> Table<E>.createEntity(row: QueryRowSet): E {
 @Suppress("UNCHECKED_CAST")
 fun <E : Entity<E>> Table<E>.createEntityWithoutReferences(row: QueryRowSet): E {
     val entity = doCreateEntity(row, skipReferences = true) as E
-    return entity.apply { discardChanges() }
+    return entity.apply { clearChangesRecursively() }
 }
 
 private fun Table<*>.doCreateEntity(row: QueryRowSet, skipReferences: Boolean = false): Entity<*> {
@@ -154,12 +154,12 @@ private fun QueryRowSet.retrieveColumn(column: Column<*>, intoEntity: Entity<*>,
                 skipReferences -> {
                     val child = Entity.create(binding.onProperty.returnType.classifier as KClass<*>, fromTable = rightTable)
                     child.implementation.setColumnValue(primaryKey, columnValue)
-                    intoEntity[binding.onProperty.name] = child.apply { discardChanges() }
+                    intoEntity[binding.onProperty.name] = child
                 }
                 this.hasColumn(primaryKey) && this[primaryKey] != null -> {
                     val child = rightTable.doCreateEntity(this)
                     child.implementation.setColumnValue(primaryKey, columnValue, forceSet = true)
-                    intoEntity[binding.onProperty.name] = child.apply { discardChanges() }
+                    intoEntity[binding.onProperty.name] = child
                 }
             }
         }

ktorm-core/src/test/kotlin/me/liuwj/ktorm/entity/EntityTest.kt

@@ -2,9 +2,7 @@ package me.liuwj.ktorm.entity
 
 import me.liuwj.ktorm.BaseTest
 import me.liuwj.ktorm.dsl.*
-import me.liuwj.ktorm.schema.Table
-import me.liuwj.ktorm.schema.int
-import me.liuwj.ktorm.schema.varchar
+import me.liuwj.ktorm.schema.*
 import org.junit.Test
 import java.io.ByteArrayInputStream
 import java.io.ByteArrayOutputStream
@@ -286,16 +284,23 @@ class EntityTest : BaseTest() {
     }
 
     interface Emp : Entity<Emp> {
+        companion object : Entity.Factory<Emp>()
         val id: Int
         var employee: Employee
         var manager: Employee
+        var hireDate: LocalDate
+        var salary: Long
+        var departmentId: Int
     }
 
     object Emps : Table<Emp>("t_employee") {
         val id by int("id").primaryKey().bindTo { it.id }
         val name by varchar("name").bindTo { it.employee.name }
         val job by varchar("job").bindTo { it.employee.job }
         val managerId by int("manager_id").bindTo { it.manager.id }
+        val hireDate by date("hire_date").bindTo { it.hireDate }
+        val salary by long("salary").bindTo { it.salary }
+        val departmentId by int("department_id").bindTo { it.departmentId }
     }
 
     @Test
@@ -304,6 +309,28 @@ class EntityTest : BaseTest() {
         emp1.employee.name = "jerry"
         // emp1.flushChanges()
 
+        val emp2 = Emp {
+            employee = emp1.employee
+            hireDate = LocalDate.now()
+            salary = 100
+            departmentId = 1
+        }
+
+        try {
+            Emps.add(emp2)
+            throw AssertionError("failed")
+
+        } catch (e: IllegalStateException) {
+            assert(e.message == "this.employee.name may be unexpectedly discarded, please save it to database first.")
+        }
+    }
+
+    @Test
+    fun testCheckUnexpectedFlush0() {
+        val emp1 = Emps.findById(1) ?: return
+        emp1.employee.name = "jerry"
+        // emp1.flushChanges()
+
         val emp2 = Emps.findById(2) ?: return
         emp2.employee = emp1.employee
 
@@ -312,7 +339,7 @@ class EntityTest : BaseTest() {
             throw AssertionError("failed")
 
         } catch (e: IllegalStateException) {
-            assert(e.message == "this.employee.name may be unexpectedly discarded after flushChanges, please save it to database first.")
+            assert(e.message == "this.employee.name may be unexpectedly discarded, please save it to database first.")
         }
     }
 
@@ -330,7 +357,7 @@ class EntityTest : BaseTest() {
             throw AssertionError("failed")
 
         } catch (e: IllegalStateException) {
-            assert(e.message == "this.employee.name may be unexpectedly discarded after flushChanges, please save it to database first.")
+            assert(e.message == "this.employee.name may be unexpectedly discarded, please save it to database first.")
         }
     }