Add support for union all Kotlin queries

Author: jeffgbutler

src/main/kotlin/org/mybatis/dynamic/sql/util/kotlin/KotlinQueryBuilder.kt

@@ -52,7 +52,12 @@ class KotlinQueryBuilder(private val dsl: QueryExpressionDSL<SelectModel>) :
 
     fun union(union: KotlinUnionBuilder.() -> QueryExpressionDSL<SelectModel>) =
         apply {
-            union(KotlinUnionBuilder(dsl))
+            union(KotlinUnionBuilder(dsl, dsl.union()))
+        }
+
+    fun unionAll(unionAll: KotlinUnionBuilder.() -> QueryExpressionDSL<SelectModel>) =
+        apply {
+            unionAll(KotlinUnionBuilder(dsl, dsl.unionAll()))
         }
 
     override fun build(): SelectModel = dsl.build()

src/main/kotlin/org/mybatis/dynamic/sql/util/kotlin/KotlinUnionBuilders.kt

@@ -20,22 +20,25 @@ import org.mybatis.dynamic.sql.SqlTable
 import org.mybatis.dynamic.sql.select.QueryExpressionDSL
 import org.mybatis.dynamic.sql.select.SelectModel
 
-class KotlinUnionBuilder(private val dsl: QueryExpressionDSL<SelectModel>) {
+class KotlinUnionBuilder(
+    private val outerDsl: QueryExpressionDSL<SelectModel>,
+    private val unionBuilder: QueryExpressionDSL<SelectModel>.UnionBuilder
+) {
     fun select(vararg selectList: BasicColumn) =
         select(selectList.toList())
 
     fun select(selectList: List<BasicColumn>) =
-        KotlinUnionFromGatherer(dsl, dsl.union().select(selectList))
+        KotlinUnionFromGatherer(outerDsl, unionBuilder.select(selectList))
 
     fun selectDistinct(vararg selectList: BasicColumn) =
         selectDistinct(selectList.toList())
 
     fun selectDistinct(selectList: List<BasicColumn>) =
-        KotlinUnionFromGatherer(dsl, dsl.union().selectDistinct(selectList))
+        KotlinUnionFromGatherer(outerDsl, unionBuilder.selectDistinct(selectList))
 }
 
 class KotlinUnionFromGatherer(
-    private val dsl: QueryExpressionDSL<SelectModel>,
+    private val outerDsl: QueryExpressionDSL<SelectModel>,
     private val fromGatherer: QueryExpressionDSL.FromGatherer<SelectModel>
 ) {
     fun from(
@@ -44,7 +47,7 @@ class KotlinUnionFromGatherer(
     ): QueryExpressionDSL<SelectModel> {
         val unionBuilder = KotlinUnionQueryBuilder(fromGatherer.from(table))
         enhance(unionBuilder)
-        return dsl
+        return outerDsl
     }
 
     fun from(
@@ -54,14 +57,14 @@ class KotlinUnionFromGatherer(
     ): QueryExpressionDSL<SelectModel> {
         val unionBuilder = KotlinUnionQueryBuilder(fromGatherer.from(table, alias))
         enhance(unionBuilder)
-        return dsl
+        return outerDsl
     }
 }
 
-class KotlinUnionQueryBuilder(private val dsl: QueryExpressionDSL<SelectModel>) :
+class KotlinUnionQueryBuilder(private val unionDsl: QueryExpressionDSL<SelectModel>) :
     KotlinBaseJoiningBuilder<QueryExpressionDSL<SelectModel>, QueryExpressionDSL<SelectModel>.QueryExpressionWhereBuilder,
-            KotlinUnionQueryBuilder>(dsl) {
+            KotlinUnionQueryBuilder>(unionDsl) {
     override fun self() = this
 
-    override fun getWhere(): QueryExpressionDSL<SelectModel>.QueryExpressionWhereBuilder = dsl.where()
+    override fun getWhere(): QueryExpressionDSL<SelectModel>.QueryExpressionWhereBuilder = unionDsl.where()
 }

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

@@ -412,7 +412,7 @@ class CanonicalSpringKotlinTest {
         }
 
         assertThat(records).hasSize(3)
-        with(records[0]!!) {
+        with(records[0]) {
             assertThat(id).isEqualTo(1)
             assertThat(firstName).isEqualTo("Fred")
             assertThat(lastName).isEqualTo("Flintstone")
@@ -422,7 +422,7 @@ class CanonicalSpringKotlinTest {
             assertThat(addressId).isEqualTo(1)
         }
 
-        with(records[2]!!) {
+        with(records[2]) {
             assertThat(id).isEqualTo(3)
             assertThat(firstName).isEqualTo("Pebbles")
             assertThat(lastName).isEqualTo("Flintstone")
@@ -485,7 +485,7 @@ class CanonicalSpringKotlinTest {
         }
 
         assertThat(records).hasSize(3)
-        with(records[0]!!) {
+        with(records[0]) {
             assertThat(id).isEqualTo(1)
             assertThat(firstName).isEqualTo("Fred")
             assertThat(lastName).isEqualTo("Flintstone")
@@ -495,7 +495,7 @@ class CanonicalSpringKotlinTest {
             assertThat(addressId).isEqualTo(1)
         }
 
-        with(records[2]!!) {
+        with(records[2]) {
             assertThat(id).isEqualTo(3)
             assertThat(firstName).isEqualTo("Pebbles")
             assertThat(lastName).isEqualTo("Flintstone")
@@ -558,7 +558,7 @@ class CanonicalSpringKotlinTest {
         }
 
         assertThat(records).hasSize(3)
-        with(records[0]!!) {
+        with(records[0]) {
             assertThat(id).isEqualTo(1)
             assertThat(firstName).isEqualTo("Fred")
             assertThat(lastName).isEqualTo("Flintstone")
@@ -568,7 +568,80 @@ class CanonicalSpringKotlinTest {
             assertThat(addressId).isEqualTo(1)
         }
 
-        with(records[2]!!) {
+        with(records[2]) {
+            assertThat(id).isEqualTo(3)
+            assertThat(firstName).isEqualTo("Pebbles")
+            assertThat(lastName).isEqualTo("Flintstone")
+            assertThat(birthDate).isNotNull()
+            assertThat(employed).isEqualTo("No")
+            assertThat(occupation).isNull()
+            assertThat(addressId).isEqualTo(1)
+        }
+    }
+
+    @Test
+    fun testRawSelectWithUnionAllAndDistinct() {
+        val selectStatement = select(
+            id.`as`("A_ID"), firstName, lastName, birthDate, employed, occupation,
+            addressId
+        ).from(Person) {
+            where(id, isEqualTo(1))
+            union {
+                select(
+                    id.`as`("A_ID"), firstName, lastName, birthDate, employed, occupation,
+                    addressId
+                ).from(Person) {
+                    where(id, isEqualTo(2))
+                }
+            }
+            unionAll {
+                selectDistinct(
+                    id.`as`("A_ID"), firstName, lastName, birthDate, employed, occupation,
+                    addressId
+                ).from(Person, "p") {
+                    where(id, isEqualTo(3))
+                }
+            }
+        }
+
+        val expected = "select id as A_ID, first_name, last_name, birth_date, employed, occupation, address_id " +
+                "from Person " +
+                "where id = :p1 " +
+                "union " +
+                "select id as A_ID, first_name, last_name, birth_date, employed, occupation, address_id " +
+                "from Person " +
+                "where id = :p2 " +
+                "union all " +
+                "select distinct p.id as A_ID, p.first_name, p.last_name, p.birth_date, p.employed, p.occupation, p.address_id " +
+                "from Person p " +
+                "where p.id = :p3"
+
+        assertThat(selectStatement.selectStatement).isEqualTo(expected)
+
+        val records = template.selectList(selectStatement) { rs, _ ->
+            val record = PersonRecord()
+            record.id = rs.getInt(1)
+            record.firstName = rs.getString(2)
+            record.lastName = rs.getString(3)
+            record.birthDate = rs.getTimestamp(4)
+            record.employed = rs.getString(5)
+            record.occupation = rs.getString(6)
+            record.addressId = rs.getInt(7)
+            record
+        }
+
+        assertThat(records).hasSize(3)
+        with(records[0]) {
+            assertThat(id).isEqualTo(1)
+            assertThat(firstName).isEqualTo("Fred")
+            assertThat(lastName).isEqualTo("Flintstone")
+            assertThat(birthDate).isNotNull()
+            assertThat(employed).isEqualTo("Yes")
+            assertThat(occupation).isEqualTo("Brontosaurus Operator")
+            assertThat(addressId).isEqualTo(1)
+        }
+
+        with(records[2]) {
             assertThat(id).isEqualTo(3)
             assertThat(firstName).isEqualTo("Pebbles")
             assertThat(lastName).isEqualTo("Flintstone")