Add KeyHolder support to the template direct inserts

Author: jeffgbutler

src/main/kotlin/org/mybatis/dynamic/sql/util/kotlin/spring/NamedParameterJdbcTemplateExtensions.kt

@@ -114,6 +114,10 @@ fun <T> NamedParameterJdbcTemplate.insertMultiple(insertStatement: MultiRowInser
                                                   keyHolder: KeyHolder) =
     update(insertStatement.insertStatement, BeanPropertySqlParameterSource(insertStatement), keyHolder)
 
+// insert with KeyHolder support
+fun NamedParameterJdbcTemplate.withKeyHolder(keyHolder: KeyHolder) =
+    KeyHolderHelper(keyHolder, this)
+
 fun NamedParameterJdbcTemplate.select(vararg selectList: BasicColumn) =
     SelectListFromGatherer(selectList.toList(), this)
 
@@ -211,20 +215,39 @@ class SelectOneMapperGatherer(
         template.selectOne(selectStatement, rowMapper)
 }
 
+@MyBatisDslMarker
+class KeyHolderHelper(private val keyHolder: KeyHolder, private val template: NamedParameterJdbcTemplate) {
+    fun insertInto(table: SqlTable, completer: GeneralInsertCompleter) =
+        template.insert(org.mybatis.dynamic.sql.util.kotlin.spring.insertInto(table, completer), keyHolder)
+
+    fun <T> insert(record: T) =
+        SingleRowInsertHelper(record, template, keyHolder)
+
+    fun <T> insertMultiple(vararg records: T) =
+        insertMultiple(records.asList())
+
+    fun <T> insertMultiple(records: List<T>) =
+        MultiRowInsertHelper(records, template, keyHolder)
+}
+
 @MyBatisDslMarker
 class BatchInsertHelper<T>(private val records: List<T>, private val template: NamedParameterJdbcTemplate) {
     fun into(table: SqlTable, completer: BatchInsertCompleter<T>) =
         template.insert(SqlBuilder.insert(records).into(table, completer))
 }
 
 @MyBatisDslMarker
-class MultiRowInsertHelper<T>(private val records: List<T>, private val template: NamedParameterJdbcTemplate) {
+class MultiRowInsertHelper<T>(private val records: List<T>, private val template: NamedParameterJdbcTemplate,
+                              private val keyHolder: KeyHolder? = null) {
     fun into(table: SqlTable, completer: MultiRowInsertCompleter<T>) =
-        template.insertMultiple(SqlBuilder.insertMultiple(records).into(table, completer))
+        keyHolder?.let { template.insertMultiple(SqlBuilder.insertMultiple(records).into(table, completer), it) }
+            ?:template.insertMultiple(SqlBuilder.insertMultiple(records).into(table, completer))
 }
 
 @MyBatisDslMarker
-class SingleRowInsertHelper<T>(private val record: T, private val template: NamedParameterJdbcTemplate) {
+class SingleRowInsertHelper<T>(private val record: T, private val template: NamedParameterJdbcTemplate,
+                               private val keyHolder: KeyHolder? = null) {
     fun into(table: SqlTable, completer: InsertCompleter<T>) =
-        template.insert(SqlBuilder.insert(record).into(table, completer))
+        keyHolder?.let { template.insert(SqlBuilder.insert(record).into(table, completer), it) }
+            ?:template.insert(SqlBuilder.insert(record).into(table, completer))
 }

src/test/kotlin/examples/kotlin/spring/canonical/CanonicalSpringKotlinTemplateDirectTest.kt

@@ -16,6 +16,7 @@
 package examples.kotlin.spring.canonical
 
 import examples.kotlin.spring.canonical.AddressDynamicSqlSupport.Address
+import examples.kotlin.spring.canonical.GeneratedAlwaysDynamicSqlSupport.GeneratedAlways
 import examples.kotlin.spring.canonical.PersonDynamicSqlSupport.Person
 import examples.kotlin.spring.canonical.PersonDynamicSqlSupport.Person.addressId
 import examples.kotlin.spring.canonical.PersonDynamicSqlSupport.Person.birthDate
@@ -32,6 +33,7 @@ import org.mybatis.dynamic.sql.util.kotlin.spring.*
 import org.springframework.jdbc.core.namedparam.NamedParameterJdbcTemplate
 import org.springframework.jdbc.datasource.embedded.EmbeddedDatabaseBuilder
 import org.springframework.jdbc.datasource.embedded.EmbeddedDatabaseType
+import org.springframework.jdbc.support.GeneratedKeyHolder
 import java.util.*
 
 class CanonicalSpringKotlinTemplateDirectTest {
@@ -246,16 +248,54 @@ class CanonicalSpringKotlinTemplateDirectTest {
         assertThat(rows[1]).isEqualTo(1)
     }
 
+    @Test
     fun testGeneralInsertWithGeneratedKey() {
-        TODO()
+        val keyHolder = GeneratedKeyHolder()
+
+        val rows = template.withKeyHolder(keyHolder).insertInto(GeneratedAlways) {
+            set(GeneratedAlways.firstName).toValue("Fred")
+            set(GeneratedAlways.lastName).toValue("Flintstone")
+        }
+
+        assertThat(rows).isEqualTo(1)
+        assertThat(keyHolder.keys).containsEntry("ID", 22)
+        assertThat(keyHolder.keys).containsEntry("FULL_NAME", "Fred Flintstone")
     }
 
+    @Test
     fun testInsertWithGeneratedKey() {
-        TODO()
+        val record = GeneratedAlwaysRecord(firstName = "Fred", lastName = "Flintstone")
+
+        val keyHolder = GeneratedKeyHolder()
+
+        val rows = template.withKeyHolder(keyHolder).insert(record).into(GeneratedAlways) {
+            map(GeneratedAlways.firstName).toProperty("firstName")
+            map(GeneratedAlways.lastName).toProperty("lastName")
+        }
+
+        assertThat(rows).isEqualTo(1)
+        assertThat(keyHolder.keys).containsEntry("ID", 22)
+        assertThat(keyHolder.keys).containsEntry("FULL_NAME", "Fred Flintstone")
     }
 
+    @Test
     fun testMultiRowInsertWithGeneratedKey() {
-        TODO()
+        val record1 = GeneratedAlwaysRecord(firstName = "Fred", lastName = "Flintstone")
+        val record2 = GeneratedAlwaysRecord(firstName = "Barney", lastName = "Rubble")
+
+        val keyHolder = GeneratedKeyHolder()
+
+        val rows = template.withKeyHolder(keyHolder).insertMultiple(record1, record2)
+            .into(GeneratedAlways) {
+                map(GeneratedAlways.firstName).toProperty("firstName")
+                map(GeneratedAlways.lastName).toProperty("lastName")
+            }
+
+        assertThat(rows).isEqualTo(2)
+        assertThat(keyHolder.keyList[0]).containsEntry("ID", 22)
+        assertThat(keyHolder.keyList[0]).containsEntry("FULL_NAME", "Fred Flintstone")
+        assertThat(keyHolder.keyList[1]).containsEntry("ID", 23)
+        assertThat(keyHolder.keyList[1]).containsEntry("FULL_NAME", "Barney Rubble")
     }
 
     @Test

src/test/kotlin/examples/kotlin/spring/canonical/CanonicalSpringKotlinTest.kt

@@ -355,8 +355,8 @@ class CanonicalSpringKotlinTest {
         val record = GeneratedAlwaysRecord(firstName = "Fred", lastName = "Flintstone")
 
         val insertStatement = insert(record).into(GeneratedAlways) {
-            map(firstName).toProperty("firstName")
-            map(lastName).toProperty("lastName")
+            map(GeneratedAlways.firstName).toProperty("firstName")
+            map(GeneratedAlways.lastName).toProperty("lastName")
         }
 
         val keyHolder = GeneratedKeyHolder()
@@ -374,8 +374,8 @@ class CanonicalSpringKotlinTest {
 
         val insertStatement = insertMultiple(record1, record2)
             .into(GeneratedAlways) {
-            map(firstName).toProperty("firstName")
-            map(lastName).toProperty("lastName")
+            map(GeneratedAlways.firstName).toProperty("firstName")
+            map(GeneratedAlways.lastName).toProperty("lastName")
         }
 
         val keyHolder = GeneratedKeyHolder()