Translate the first branch of a UNION query

core/src/main/kotlin/org/evomaster/core/database/sql/solver/SmtLibGenerator.kt

@@ -4,8 +4,11 @@ import net.sf.jsqlparser.schema.Table
 import net.sf.jsqlparser.statement.Statement
 import net.sf.jsqlparser.statement.delete.Delete
 import net.sf.jsqlparser.statement.select.FromItem
+import net.sf.jsqlparser.statement.select.ParenthesedSelect
 import net.sf.jsqlparser.statement.select.PlainSelect
 import net.sf.jsqlparser.statement.select.Select
+import net.sf.jsqlparser.statement.select.SetOperationList
+import net.sf.jsqlparser.statement.select.UnionOp
 import net.sf.jsqlparser.statement.update.Update
 import net.sf.jsqlparser.util.TablesNamesFinder
 import org.evomaster.client.java.controller.api.dto.database.schema.DatabaseType
@@ -554,8 +557,8 @@ class SmtLibGenerator(
 
         val (where, defaultTable) = when (sqlQuery) {
             is Select -> {
-                val plainSelect = sqlQuery.selectBody as PlainSelect
-                Pair(plainSelect.where, defaultTableOf(sqlQuery, plainSelect))
+                val plainSelect = plainSelectOf(sqlQuery)
+                Pair(plainSelect.where, defaultTableOf(plainSelect))
             }
             is Delete -> Pair(sqlQuery.where, sqlQuery.table.getName())
             is Update -> Pair(sqlQuery.where, sqlQuery.table.getName())
@@ -654,18 +657,20 @@ class SmtLibGenerator(
         queryTables: Collection<String>
     ) {
         if (sqlQuery is Select) { // TODO: Handle other queries
-            val plainSelect = sqlQuery.selectBody as PlainSelect
+            val plainSelect = plainSelectOf(sqlQuery)
             val joins = plainSelect.joins
             if (joins != null) {
                 for (join in joins) {
                     val onExpressions = join.onExpressions
                     if (onExpressions.isNotEmpty()) {
-                        // TODO: only the first ON expression is used; a composite ON
-                        // (e.g. "a = b AND c = d") drops all but the first conjunct.
+                        // TODO: only the first ON expression is used. A join has more than one only in
+                        // the nested form "a JOIN b JOIN c ON b.x = c.x ON a.x = b.x", where every ON after
+                        // the first is dropped. A composite ON such as "a = b AND c = d" is a single
+                        // expression, and each of its conjuncts is translated.
                         val onExpression = onExpressions.elementAt(0)
                         try {
                             val condition = parser.parse(onExpression.toString(), toDBType(schema.databaseType))
-                            val tableFromQuery = defaultTableOf(sqlQuery, plainSelect)!!
+                            val tableFromQuery = defaultTableOf(plainSelect)!!
                             // TODO: the ON condition is translated with the SAME row index on
                             // both sides ("diagonal pairing"): row i of one table is matched only with row i
                             // of the other. This is sufficient at the default numberOfRows=1 to force a
@@ -717,10 +722,13 @@ class SmtLibGenerator(
      * The table an unqualified column falls back to when no single query table declares it: the
      * table in `FROM`. [TablesNamesFinder] returns its tables in a hash set, so its first element is
      * not necessarily that table — in `FROM products p JOIN users u` it is `users`.
+     *
+     * [plainSelect] is the SELECT being translated, the first branch of a UNION, so the fallback only
+     * looks at the tables of that branch.
      */
-    private fun defaultTableOf(sqlQuery: Select, plainSelect: PlainSelect): String? =
+    private fun defaultTableOf(plainSelect: PlainSelect): String? =
         (plainSelect.fromItem as? Table)?.name
-            ?: TablesNamesFinder().getTables(sqlQuery as Statement).firstOrNull()
+            ?: TablesNamesFinder().getTables(plainSelect as Statement).firstOrNull()
 
     /**
      * Builds the [SMTConditionVisitor.ColumnScope] for one query from the aliases its `FROM` and
@@ -761,6 +769,27 @@ class SmtLibGenerator(
         val derivedAliases: Set<String>
     )
 
+    /**
+     * The plain `SELECT` whose clauses are translated: the query itself when it is a [PlainSelect], or
+     * the first branch of a `UNION`. Always a [PlainSelect].
+     *
+     * A `UNION` returns rows as soon as any one of its branches does, so generating data for its first
+     * branch is enough; the branch is translated as if it were the whole query. `INTERSECT` and
+     * `EXCEPT` constrain every branch, which this translation cannot express, so they are rejected
+     * and the query is recorded as untranslatable.
+     */
+    private fun plainSelectOf(select: Select): PlainSelect = when (select) {
+        is PlainSelect -> select
+        is ParenthesedSelect -> plainSelectOf(select.select)
+        is SetOperationList ->
+            if (select.operations.all { it is UnionOp }) {
+                plainSelectOf(select.selects.first())
+            } else {
+                throw RuntimeException("Unsupported set operation, only UNION can be translated: $select")
+            }
+        else -> throw RuntimeException("Unsupported SELECT form ${select.javaClass.simpleName}: $select")
+    }
+
     /**
      * Reads the aliases a query declares.
      *
@@ -774,7 +803,7 @@ class SmtLibGenerator(
 
         when (sqlQuery) {
             is Select -> {
-                val plainSelect = sqlQuery.selectBody as PlainSelect
+                val plainSelect = plainSelectOf(sqlQuery)
                 val fromItem = plainSelect.fromItem
                 if (fromItem != null) {
                     val tableName = getTableName(fromItem)
@@ -854,7 +883,7 @@ class SmtLibGenerator(
 
         // Add tables from JOINs and WHERE clause if they exist
         if (sqlQuery is Select) {
-            val plainSelect = sqlQuery.selectBody as PlainSelect
+            val plainSelect = plainSelectOf(sqlQuery)
 
             // Add tables from JOINs
             plainSelect.joins?.forEach { join ->

core/src/test/kotlin/org/evomaster/core/database/sql/solver/service/SetOperationSolvingTest.kt

@@ -0,0 +1,80 @@
+package org.evomaster.core.database.sql.solver.service
+
+import org.evomaster.client.java.controller.api.dto.database.schema.DbInfoDto
+import org.evomaster.client.java.sql.DbInfoExtractor
+import org.evomaster.client.java.sql.SqlScriptRunner
+import org.evomaster.core.database.sql.SqlActionTransformer
+import org.junit.jupiter.api.AfterAll
+import org.junit.jupiter.api.Assertions.assertFalse
+import org.junit.jupiter.api.Assertions.assertTrue
+import org.junit.jupiter.api.BeforeAll
+import org.junit.jupiter.api.Test
+import java.sql.Connection
+import java.sql.DriverManager
+
+/**
+ * A query combining SELECTs with UNION is parsed into a set-operation node rather than a plain SELECT.
+ * The generator used to cast it to a plain SELECT, so the whole query failed with a
+ * ClassCastException and no data was generated for it.
+ */
+class SetOperationSolvingTest {
+
+    companion object {
+        private lateinit var solver: SMTLibZ3DbConstraintSolver
+        private lateinit var connection: Connection
+        private lateinit var schemaDto: DbInfoDto
+
+        @JvmStatic
+        @BeforeAll
+        fun setup() {
+            connection = DriverManager.getConnection("jdbc:h2:mem:set_operation_test", "sa", "")
+            SqlScriptRunner.execCommand(
+                connection,
+                "CREATE TABLE users(id bigint primary key, age int);\n" +
+                    "CREATE TABLE products(id bigint primary key, stock int);\n"
+            )
+            schemaDto = DbInfoExtractor.extract(connection)
+            solver = SMTLibZ3DbConstraintSolver()
+            solver.initializeExecutor()
+        }
+
+        @JvmStatic
+        @AfterAll
+        fun tearDown() {
+            connection.close()
+            if (this::solver.isInitialized) {
+                solver.close()
+            }
+        }
+    }
+
+    /**
+     * A UNION returns rows as soon as one branch does, so data satisfying the first branch is enough.
+     */
+    @Test
+    fun unionGeneratesDataForItsFirstBranch() {
+        val query = "SELECT id FROM users WHERE age > 30 UNION SELECT id FROM products WHERE stock < 0"
+
+        val actions = solver.solve(schemaDto, query, 1)
+        assertFalse(actions.isEmpty(), "Solver should return actions for a satisfiable UNION")
+
+        SqlScriptRunner.execInsert(connection, SqlActionTransformer.transform(actions).insertions)
+        val result = SqlScriptRunner.execCommand(connection, query)
+        assertFalse(result.isEmpty(), "Inserted rows should satisfy the original UNION query")
+    }
+
+    /**
+     * INTERSECT constrains every branch, which the translation cannot express. It must be reported as
+     * untranslatable, not crash.
+     */
+    @Test
+    fun intersectIsRejectedWithoutCrashing() {
+        val newActions = solver.solve(
+            schemaDto,
+            "SELECT id FROM users WHERE age > 30 INTERSECT SELECT id FROM products WHERE stock < 0",
+            1
+        )
+
+        assertTrue(newActions.isEmpty())
+    }
+}