fix: Change IdTable idColumns to read-only and add validation on acce…

…ss (#2178)

* fix: Change IdTable idColumns to read-only and add validation on access

- Previous validation is removed to only check if idColumns is empty if it is accessed
- idColumns is now publicly read-only
- addIdColumn() is now protected

exposed-core/api/exposed-core.api

@@ -6,6 +6,7 @@ public final class org/jetbrains/exposed/dao/id/CompositeID : java/lang/Comparab
 	public fun equals (Ljava/lang/Object;)Z
 	public final fun get (Lorg/jetbrains/exposed/sql/Column;)Ljava/lang/Comparable;
 	public fun hashCode ()I
+	public final fun setWithEntityID (Lorg/jetbrains/exposed/sql/Column;Lorg/jetbrains/exposed/dao/id/EntityID;)V
 	public final fun setWithEntityIdValue (Lorg/jetbrains/exposed/sql/Column;Ljava/lang/Comparable;)V
 	public final fun setWithNullableEntityIdValue (Lorg/jetbrains/exposed/sql/Column;Ljava/lang/Comparable;)V
 	public fun toString ()Ljava/lang/String;
@@ -50,8 +51,9 @@ public abstract class org/jetbrains/exposed/dao/id/IdTable : org/jetbrains/expos
 	public fun <init> ()V
 	public fun <init> (Ljava/lang/String;)V
 	public synthetic fun <init> (Ljava/lang/String;ILkotlin/jvm/internal/DefaultConstructorMarker;)V
+	protected final fun addIdColumn (Lorg/jetbrains/exposed/sql/Column;)V
 	public abstract fun getId ()Lorg/jetbrains/exposed/sql/Column;
-	public final fun getIdColumns ()Ljava/util/HashSet;
+	public final fun getIdColumns ()Ljava/util/Set;
 }
 
 public class org/jetbrains/exposed/dao/id/IntIdTable : org/jetbrains/exposed/dao/id/IdTable {

exposed-core/src/main/kotlin/org/jetbrains/exposed/dao/id/CompositeID.kt

@@ -24,6 +24,14 @@ class CompositeID private constructor() : Comparable<CompositeID> {
         values[column] = value?.let { EntityID(value, column.table as IdTable<T>) }
     }
 
+    @JvmName("setWithEntityID")
+    operator fun <T : Comparable<T>, ID : EntityID<T>> set(column: Column<ID>, value: ID) {
+        require(values.isEmpty() || values.keys.first().table == column.table) {
+            "CompositeID key columns must all come from the same IdTable ${values.keys.first().table.tableName}"
+        }
+        values[column] = value
+    }
+
     @Suppress("UNCHECKED_CAST")
     operator fun <T : Comparable<T>> get(column: Column<T>): T = values[column] as T
 

exposed-core/src/main/kotlin/org/jetbrains/exposed/dao/id/IdTable.kt

@@ -35,23 +35,28 @@ abstract class IdTable<T : Comparable<T>>(name: String = "") : Table(name) {
     /** The identity column of this [IdTable], for storing values of type [T] wrapped as [EntityID] instances. */
     abstract val id: Column<EntityID<T>>
 
+    private val _idColumns = HashSet<Column<out Comparable<*>>>()
+
     /** All base columns that make up this [IdTable]'s identifier column. */
-    val idColumns = HashSet<Column<out Comparable<*>>>()
+    val idColumns: Set<Column<out Comparable<*>>>
+        get() = _idColumns.ifEmpty {
+            val message = "Table definition must include id columns. Please use Column.entityId() or IdTable.addIdColumn()."
+            exposedLogger.error(message)
+            error(message)
+        }
 
-    /**
-     * Adds a column to [idColumns].
-     *
-     * @throws IllegalStateException If more than one column is added to an [IdTable] that is not a [CompositeIdTable].
-     */
-    internal fun <S : Comparable<S>> addIdColumn(newColumn: Column<EntityID<S>>) {
-        if (idColumns.isNotEmpty() && this !is CompositeIdTable) {
+    /** Adds a column to [idColumns] so that it can be used as a component of the [id] property. */
+    protected fun <S : Comparable<S>> addIdColumn(newColumn: Column<EntityID<S>>) {
+        if (_idColumns.isNotEmpty() && this !is CompositeIdTable) {
             val message = "CompositeIdTable should be used if multiple EntityID key columns are required"
             exposedLogger.error(message)
             error(message)
         }
-        idColumns.add(newColumn)
+        _idColumns.add(newColumn)
     }
 
+    internal fun <S : Comparable<S>> addIdColumnInternal(newColumn: Column<EntityID<S>>) { addIdColumn(newColumn) }
+
     /**
      * Returns a boolean operator comparing each of this table's [idColumns] to its corresponding
      * value in [toCompare], using the specified SQL [operator].

exposed-core/src/main/kotlin/org/jetbrains/exposed/sql/Table.kt

@@ -618,7 +618,7 @@ open class Table(name: String = "") : ColumnSet(), DdlAware {
             it.dbDefaultValue = dbDefaultValue?.let { it as Expression<EntityID<T>> }
             it.extraDefinitions = extraDefinitions
         }
-        (table as IdTable<T>).addIdColumn(newColumn)
+        (table as IdTable<T>).addIdColumnInternal(newColumn)
         return replaceColumn(this, newColumn)
     }
 

exposed-tests/src/test/kotlin/org/jetbrains/exposed/sql/tests/shared/entities/CompositeIdTableEntityTest.kt

@@ -5,7 +5,12 @@ import org.jetbrains.exposed.dao.id.CompositeID
 import org.jetbrains.exposed.dao.id.CompositeIdTable
 import org.jetbrains.exposed.dao.id.EntityID
 import org.jetbrains.exposed.dao.id.IntIdTable
-import org.jetbrains.exposed.sql.*
+import org.jetbrains.exposed.sql.SchemaUtils
+import org.jetbrains.exposed.sql.and
+import org.jetbrains.exposed.sql.exists
+import org.jetbrains.exposed.sql.insert
+import org.jetbrains.exposed.sql.insertAndGetId
+import org.jetbrains.exposed.sql.selectAll
 import org.jetbrains.exposed.sql.tests.DatabaseTestsBase
 import org.jetbrains.exposed.sql.tests.TestDB
 import org.jetbrains.exposed.sql.tests.currentTestDB
@@ -108,6 +113,7 @@ class CompositeIdTableEntityTest : DatabaseTestsBase() {
         val publisherId = integer("publisher_id").nullable()
         val publisherIsbn = uuid("publisher_isbn").nullable()
 
+        override val primaryKey = PrimaryKey(zipCode, name, areaCode)
         init {
             foreignKey(publisherId, publisherIsbn, target = Publishers.primaryKey)
         }
@@ -136,6 +142,27 @@ class CompositeIdTableEntityTest : DatabaseTestsBase() {
         }
     }
 
+    @Test
+    fun testCreateWithMissingIdColumns() {
+        val missingIdsTable = object : CompositeIdTable("missing_ids_table") {
+            val age = integer("age")
+            val name = varchar("name", 50)
+            override val primaryKey = PrimaryKey(age, name)
+        }
+
+        withDb {
+            // table can be created with no issue
+            SchemaUtils.create(missingIdsTable)
+
+            expectException<IllegalStateException> {
+                // but trying to use id property requires idColumns not being empty
+                missingIdsTable.select(missingIdsTable.id).toList()
+            }
+
+            SchemaUtils.drop(missingIdsTable)
+        }
+    }
+
     @Test
     fun testInsertAndSelectUsingDAO() {
         withTables(excludeSettings = listOf(TestDB.SQLITE), Publishers) {
@@ -347,6 +374,7 @@ class CompositeIdTableEntityTest : DatabaseTestsBase() {
         val zipCode = varchar("zip_code", 8).entityId()
         val name = varchar("name", 64).entityId()
         val population = long("population").nullable()
+        override val primaryKey = PrimaryKey(zipCode, name)
     }
 
     class Town(id: EntityID<CompositeID>) : CompositeEntity(id) {

exposed-tests/src/test/kotlin/org/jetbrains/exposed/sql/tests/shared/entities/ViaTest.kt

@@ -1,6 +1,8 @@
 package org.jetbrains.exposed.sql.tests.shared.entities
 
 import org.jetbrains.exposed.dao.*
+import org.jetbrains.exposed.dao.id.CompositeID
+import org.jetbrains.exposed.dao.id.CompositeIdTable
 import org.jetbrains.exposed.dao.id.EntityID
 import org.jetbrains.exposed.dao.id.IdTable
 import org.jetbrains.exposed.dao.id.IntIdTable
@@ -10,9 +12,14 @@ import org.jetbrains.exposed.sql.tests.DatabaseTestsBase
 import org.jetbrains.exposed.sql.tests.shared.assertEqualCollections
 import org.jetbrains.exposed.sql.tests.shared.assertEqualLists
 import org.jetbrains.exposed.sql.tests.shared.assertEquals
+import org.jetbrains.exposed.sql.transactions.TransactionManager
+import org.jetbrains.exposed.sql.transactions.inTopLevelTransaction
 import org.junit.Test
+import java.sql.Connection
 import java.util.*
 import kotlin.reflect.jvm.isAccessible
+import kotlin.test.assertFalse
+import kotlin.test.assertTrue
 
 object ViaTestData {
     object NumbersTable : UUIDTable() {
@@ -280,4 +287,88 @@ class ViaTests : DatabaseTestsBase() {
             }
         }
     }
+
+    object Projects : IntIdTable("projects") {
+        val name = varchar("name", 50)
+    }
+    class Project(id: EntityID<Int>) : IntEntity(id) {
+        companion object : IntEntityClass<Project>(Projects)
+
+        var name by Projects.name
+        val tasks by Task via ProjectTasks
+    }
+
+    object ProjectTasks : CompositeIdTable("project_tasks") {
+        val project = reference("project", Projects, onDelete = ReferenceOption.CASCADE)
+        val task = reference("task", Tasks, onDelete = ReferenceOption.CASCADE)
+        val approved = bool("approved")
+
+        override val primaryKey = PrimaryKey(project, task)
+
+        init {
+            addIdColumn(project)
+            addIdColumn(task)
+        }
+    }
+    class ProjectTask(id: EntityID<CompositeID>) : CompositeEntity(id) {
+        companion object : CompositeEntityClass<ProjectTask>(ProjectTasks)
+
+        var approved by ProjectTasks.approved
+    }
+
+    object Tasks : IntIdTable("tasks") {
+        val title = varchar("title", 64)
+    }
+    class Task(id: EntityID<Int>) : IntEntity(id) {
+        companion object : IntEntityClass<Task>(Tasks)
+
+        var title by Tasks.title
+        val approved by ProjectTasks.approved
+    }
+
+    @Test
+    fun testAdditionalLinkDataUsingCompositeIdInnerTable() {
+        withTables(Projects, Tasks, ProjectTasks) {
+            val p1 = Project.new { name = "Project 1" }
+            val p2 = Project.new { name = "Project 2" }
+            val t1 = Task.new { title = "Task 1" }
+            val t2 = Task.new { title = "Task 2" }
+            val t3 = Task.new { title = "Task 3" }
+
+            ProjectTask.new(
+                CompositeID {
+                    it[ProjectTasks.task] = t1.id
+                    it[ProjectTasks.project] = p1.id
+                }
+            ) { approved = true }
+            ProjectTask.new(
+                CompositeID {
+                    it[ProjectTasks.task] = t2.id
+                    it[ProjectTasks.project] = p2.id
+                }
+            ) { approved = false }
+            ProjectTask.new(
+                CompositeID {
+                    it[ProjectTasks.task] = t3.id
+                    it[ProjectTasks.project] = p2.id
+                }
+            ) { approved = false }
+
+            commit()
+
+            inTopLevelTransaction(Connection.TRANSACTION_SERIALIZABLE) {
+                maxAttempts = 1
+                Project.all().with(Project::tasks)
+                val cache = TransactionManager.current().entityCache
+
+                val p1Tasks = cache.getReferrers<Task>(p1.id, ProjectTasks.project)?.toList().orEmpty()
+                assertEqualLists(p1Tasks.map { it.id }, listOf(t1.id))
+                assertTrue { p1Tasks.all { task -> task.approved } }
+
+                val p2Tasks = cache.getReferrers<Task>(p2.id, ProjectTasks.project)?.toList().orEmpty()
+                assertEqualLists(p2Tasks.map { it.id }, listOf(t2.id, t3.id))
+                assertFalse { p1Tasks.all { task -> !task.approved } }
+            }
+        }
+    }
 }