Fix foreign key gene

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

@@ -18,6 +18,7 @@ import org.evomaster.core.search.gene.numeric.DoubleGene
 import org.evomaster.core.search.gene.numeric.IntegerGene
 import org.evomaster.core.search.gene.numeric.LongGene
 import org.evomaster.core.search.gene.placeholder.ImmutableDataHolderGene
+import org.evomaster.core.search.gene.sql.SqlForeignKeyGene
 import org.evomaster.core.search.gene.sql.SqlPrimaryKeyGene
 import org.evomaster.core.search.gene.string.StringGene
 import org.evomaster.core.search.service.Statistics
@@ -55,6 +56,13 @@ import kotlin.io.path.createDirectories
 import kotlin.io.path.exists
 import kotlin.text.equals
 
+/**
+ * A row already turned into an action, kept so that later rows can bind their foreign keys to it.
+ *
+ * @param values the value Z3 assigned to each column, keyed by the uppercase column name.
+ */
+private class InsertedRow(val tableId: TableId, val actionId: Long, val values: Map<String, String?>)
+
 /**
  * An SMT solver implementation using Z3 in a Docker container.
  * It generates the SMT problem from the database schema and the SQL query,
@@ -360,6 +368,7 @@ class SMTLibZ3DbConstraintSolver() : DbConstraintSolver {
      */
     private fun toSqlActionList(schemaDto: DbInfoDto, solution: Z3Solution): List<SqlAction> {
         val actions = mutableListOf<SqlAction>()
+        val insertedRows = mutableListOf<InsertedRow>()
 
         for (row in inInsertionOrder(schemaDto, solution.assignments)) {
             val tableName = getTableName(row.key)
@@ -370,60 +379,131 @@ class SMTLibZ3DbConstraintSolver() : DbConstraintSolver {
             val actionId = idCounter
             idCounter++
 
-            val genes = mutableListOf<Gene>()
+            val valueGenes = mutableListOf<Gene>()
+            val rawValues = mutableMapOf<String, String?>()
             for (smtColumn in columns.fields) {
                 val dbColumn = table.columns.firstOrNull {
                     convertToAscii(it.name).equals(smtColumn, ignoreCase = true)
                 }
                 val dbColumnName = dbColumn?.name ?: smtColumn
+                val columnValue = columns.getField(smtColumn)
 
-                var gene: Gene = IntegerGene(dbColumnName, 0)
-                when (val columnValue = columns.getField(smtColumn)) {
-                    is StringValue -> {
-                        gene = if (hasColumnType(schemaDto, table, dbColumnName, SmtLibGenerator.BOOLEAN_TYPE)) {
-                            BooleanGene(dbColumnName, toBoolean(columnValue.value))
-                        } else {
-                            StringGene(dbColumnName, columnValue.value)
-                        }
-                    }
-                    is LongValue -> {
-                        gene = if (hasColumnType(schemaDto, table, dbColumnName, SmtLibGenerator.TIMESTAMP_TYPE)) {
-                            val epochSeconds = columnValue.value.toLong()
-                            val localDateTime = LocalDateTime.ofInstant(
-                                Instant.ofEpochSecond(epochSeconds), ZoneOffset.UTC
-                            )
-                            val formatted = localDateTime.format(
-                                DateTimeFormatter.ofPattern(TIMESTAMP_FORMAT)
-                            )
-                            ImmutableDataHolderGene(dbColumnName, formatted, inQuotes = true)
-                        } else {
-                            LongGene(dbColumnName, columnValue.value.toLong())
-                        }
-                    }
-                    is RealValue -> {
-                        gene = DoubleGene(dbColumnName, columnValue.value)
-                    }
-                }
+                rawValues[dbColumnName.uppercase()] = rawValue(columnValue)
+                valueGenes.add(toValueGene(schemaDto, table, dbColumnName, columnValue))
+            }
+
+            val foreignKeyGenes = bindForeignKeys(table, actionId, rawValues, insertedRows)
+
+            val genes = valueGenes.map { valueGene ->
+                var gene: Gene = foreignKeyGenes[valueGene.name.uppercase()] ?: valueGene
+                val dbColumn = table.columns.firstOrNull { it.name.equals(valueGene.name, ignoreCase = true) }
                 if (dbColumn != null && dbColumn.primaryKey) {
-                    gene = SqlPrimaryKeyGene(dbColumnName, table.id, gene, actionId)
+                    gene = SqlPrimaryKeyGene(valueGene.name, table.id, gene, actionId)
                 }
-                /*
-                    TODO: a foreign key column is returned as a plain value gene, not as a
-                    SqlForeignKeyGene bound to the action of the referenced row. Z3 makes the value
-                    match that row's primary key, but once the search mutates either gene the
-                    reference can break, and FK repair does not see this column as a foreign key.
-                 */
                 gene.markAllAsInitialized()
-                genes.add(gene)
+                gene
             }
 
-            val sqlAction = SqlAction(table, table.columns, actionId, genes.toList())
-            actions.add(sqlAction)
+            actions.add(SqlAction(table, table.columns, actionId, genes))
+            insertedRows.add(InsertedRow(table.id, actionId, rawValues))
         }
 
         return actions
     }
 
+    /**
+     * Builds a [SqlForeignKeyGene] for each foreign key column of a row, bound to the action of the row
+     * it references.
+     *
+     * Z3 makes the value of a foreign key equal to the primary key of some row of the referenced
+     * table, but a plain value gene would only copy that value: once the search mutates either side,
+     * the reference breaks, and the FK repair of EvoMaster does not recognise the column as a foreign
+     * key. Binding the gene to the referenced action keeps the two linked, as for any other insertion.
+     *
+     * The referenced row is the one, among the rows inserted before this one, whose target columns hold
+     * the values Z3 assigned to the source columns. A foreign key is left as plain values when no such
+     * row exists: when its target columns are not the primary key, which a [SqlForeignKeyGene] cannot
+     * refer to, or when a row references itself, since no action precedes it.
+     *
+     * @return the foreign key genes, keyed by the uppercase name of their source column.
+     */
+    private fun bindForeignKeys(
+        table: Table,
+        actionId: Long,
+        rawValues: Map<String, String?>,
+        insertedRows: List<InsertedRow>
+    ): Map<String, SqlForeignKeyGene> {
+        val genes = mutableMapOf<String, SqlForeignKeyGene>()
+
+        for (foreignKey in table.foreignKeys) {
+            if (foreignKey.targetColumns.any { !it.primaryKey }) continue
+
+            val referenced = insertedRows.lastOrNull { row ->
+                row.tableId == foreignKey.targetTableId &&
+                    foreignKey.sourceColumns.indices.all { i ->
+                        val value = rawValues[foreignKey.sourceColumns[i].name.uppercase()]
+                        value != null && value == row.values[foreignKey.targetColumns[i].name.uppercase()]
+                    }
+            } ?: continue
+
+            foreignKey.sourceColumns.forEachIndexed { i, sourceColumn ->
+                genes[sourceColumn.name.uppercase()] = SqlForeignKeyGene(
+                    sourceColumn = sourceColumn.name,
+                    uniqueId = actionId,
+                    targetTable = foreignKey.targetTableId,
+                    targetColumn = foreignKey.targetColumns[i].name,
+                    nullable = sourceColumn.nullable,
+                    uniqueIdOfPrimaryKey = referenced.actionId,
+                    otherSourceColumnsInCompositeFK = foreignKey.sourceColumns
+                        .filter { it != sourceColumn }.map { it.name }.toSet()
+                )
+            }
+        }
+
+        return genes
+    }
+
+    /**
+     * The value Z3 assigned to a column, as a string that can be compared across rows.
+     */
+    private fun rawValue(value: SMTLibValue?): String? = when (value) {
+        is LongValue -> value.value.toString()
+        is RealValue -> value.value.toString()
+        is StringValue -> value.value
+        else -> null
+    }
+
+    /**
+     * Builds the gene holding the value Z3 assigned to a column, according to the column's SQL type.
+     */
+    private fun toValueGene(schemaDto: DbInfoDto, table: Table, dbColumnName: String, columnValue: SMTLibValue?): Gene {
+        return when (columnValue) {
+            is StringValue -> {
+                if (hasColumnType(schemaDto, table, dbColumnName, SmtLibGenerator.BOOLEAN_TYPE)) {
+                    BooleanGene(dbColumnName, toBoolean(columnValue.value))
+                } else {
+                    StringGene(dbColumnName, columnValue.value)
+                }
+            }
+            is LongValue -> {
+                if (hasColumnType(schemaDto, table, dbColumnName, SmtLibGenerator.TIMESTAMP_TYPE)) {
+                    val epochSeconds = columnValue.value.toLong()
+                    val localDateTime = LocalDateTime.ofInstant(
+                        Instant.ofEpochSecond(epochSeconds), ZoneOffset.UTC
+                    )
+                    val formatted = localDateTime.format(
+                        DateTimeFormatter.ofPattern(TIMESTAMP_FORMAT)
+                    )
+                    ImmutableDataHolderGene(dbColumnName, formatted, inQuotes = true)
+                } else {
+                    LongGene(dbColumnName, columnValue.value.toLong())
+                }
+            }
+            is RealValue -> DoubleGene(dbColumnName, columnValue.value)
+            else -> IntegerGene(dbColumnName, 0)
+        }
+    }
+
     /**
      * Orders the rows of a solution so that each one comes after the rows it references through a
      * foreign key, and rows of the same table by their index.
@@ -554,7 +634,7 @@ class SMTLibZ3DbConstraintSolver() : DbConstraintSolver {
         return Table(
             TableId.fromDto(schema.databaseType, tableDto.id),
             findColumns(schema, tableDto),
-            findForeignKeys(tableDto) // TODO: Implement this method
+            findForeignKeys(schema, tableDto)
         )
     }
 
@@ -657,9 +737,34 @@ class SMTLibZ3DbConstraintSolver() : DbConstraintSolver {
         }
     }
 
-    // TODO: Implement this method
-    private fun findForeignKeys(tableDto: TableDto): Set<ForeignKey> {
-        return emptySet()
+    /**
+     * Rebuilds the foreign keys of a table from the schema, in the same way as
+     * [org.evomaster.core.database.sql.SqlInsertBuilder]: when the schema does not name the target
+     * columns, the foreign key references the primary key of the target table. A foreign key whose
+     * table or columns cannot be resolved is skipped.
+     */
+    private fun findForeignKeys(schema: DbInfoDto, tableDto: TableDto): Set<ForeignKey> {
+        val sourceColumns = findColumns(schema, tableDto)
+
+        return tableDto.foreignKeys.mapNotNull { foreignKey ->
+            val targetDto = schema.tables.find { it.id.name.equals(foreignKey.targetTable, ignoreCase = true) }
+                ?: return@mapNotNull null
+            val targetColumns = findColumns(schema, targetDto)
+
+            val source = foreignKey.sourceColumns.map { name ->
+                sourceColumns.find { it.name.equals(name, ignoreCase = true) } ?: return@mapNotNull null
+            }
+            val target = if (foreignKey.targetColumns.isEmpty()) {
+                targetColumns.filter { it.primaryKey }
+            } else {
+                foreignKey.targetColumns.map { name ->
+                    targetColumns.find { it.name.equals(name, ignoreCase = true) } ?: return@mapNotNull null
+                }
+            }
+            if (source.isEmpty() || source.size != target.size) return@mapNotNull null
+
+            ForeignKey(source, TableId.fromDto(schema.databaseType, targetDto.id), target)
+        }.toSet()
     }
 
     /**

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

@@ -5,7 +5,12 @@ 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.evomaster.core.search.gene.numeric.LongGene
+import org.evomaster.core.search.gene.sql.SqlForeignKeyGene
+import org.evomaster.core.search.gene.sql.SqlPrimaryKeyGene
+import org.junit.jupiter.api.Assertions.assertEquals
 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.BeforeEach
 import org.junit.jupiter.api.Test
@@ -74,6 +79,32 @@ class ForeignKeyReferencedRowsTest {
         insertSolutionAndQuery("SELECT * FROM chapter WHERE number > 3", 2)
     }
 
+    /**
+     * The foreign key column is bound to the action of the referenced row, not copied as a plain value,
+     * so the reference survives when the search later mutates the referenced primary key.
+     */
+    @Test
+    fun foreignKeyStaysBoundWhenTheReferencedPrimaryKeyChanges() {
+        val actions = solver.solve(schemaDto, "SELECT * FROM book WHERE pages > 10", 1)
+
+        val author = actions.single { it.table.id.name.equals("AUTHOR", ignoreCase = true) }
+        val book = actions.single { it.table.id.name.equals("BOOK", ignoreCase = true) }
+
+        val fk = book.seeTopGenes().single { it.name.equals("AUTHOR", ignoreCase = true) }
+        assertTrue(fk is SqlForeignKeyGene) { "expected a foreign key gene, got ${fk::class.simpleName}" }
+        assertEquals(author.insertionId, (fk as SqlForeignKeyGene).uniqueIdOfPrimaryKey)
+
+        // Change the author's id, as a mutation would
+        val pk = author.seeTopGenes().single { it.name.equals("ID", ignoreCase = true) } as SqlPrimaryKeyGene
+        (pk.gene as LongGene).value = 4242
+
+        val dto = SqlActionTransformer.transform(actions)
+        SqlScriptRunner.execInsert(connection, dto.insertions)
+
+        val result = SqlScriptRunner.execCommand(connection, "SELECT * FROM book WHERE author = 4242")
+        assertFalse(result.isEmpty(), "The book should reference the author through its new id")
+    }
+
     private fun insertSolutionAndQuery(selectQuery: String, numberOfRows: Int) {
         val actions = solver.solve(schemaDto, selectQuery, numberOfRows)
         assertFalse(actions.isEmpty(), "Solver should return actions for a satisfiable query")