feat: EXPOSED-77 Support entity class for table with composite primar…

…y key

- Extract logic between regular id columns and composite id columns to lower locations.
- Add KDocs and change some object locations.
- Flesh out unit tests

exposed-core/api/exposed-core.api

@@ -1,5 +1,5 @@
 public final class org/jetbrains/exposed/dao/id/CompositeID : java/util/HashMap, java/lang/Comparable {
-	public fun <init> ()V
+	public fun <init> (Ljava/util/Map;)V
 	public synthetic fun compareTo (Ljava/lang/Object;)I
 	public fun compareTo (Lorg/jetbrains/exposed/dao/id/CompositeID;)I
 	public final fun containsKey (Ljava/lang/Object;)Z
@@ -30,11 +30,12 @@ public final class org/jetbrains/exposed/dao/id/CompositeID : java/util/HashMap,
 	public final fun values ()Ljava/util/Collection;
 }
 
-public abstract class org/jetbrains/exposed/dao/id/CompositeIdTable : org/jetbrains/exposed/dao/id/IdTable {
+public class org/jetbrains/exposed/dao/id/CompositeIdTable : org/jetbrains/exposed/dao/id/IdTable {
 	public fun <init> ()V
 	public fun <init> (Ljava/lang/String;)V
 	public synthetic fun <init> (Ljava/lang/String;ILkotlin/jvm/internal/DefaultConstructorMarker;)V
-	public final fun compositeEntityId (Lorg/jetbrains/exposed/sql/Column;[Lorg/jetbrains/exposed/sql/Column;)Lorg/jetbrains/exposed/sql/Column;
+	public final fun compositeEntityId (Lorg/jetbrains/exposed/sql/Column;)Lorg/jetbrains/exposed/sql/Column;
+	public final fun getId ()Lorg/jetbrains/exposed/sql/Column;
 	public final fun getIdColumns ()Ljava/util/HashSet;
 }
 

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

@@ -0,0 +1,33 @@
+package org.jetbrains.exposed.dao.id
+
+import org.jetbrains.exposed.sql.Column
+
+/** Class representing a mapping of each composite primary key column to its stored identity value. */
+class CompositeID(
+    idMapping: Map<out Column<*>, Comparable<*>?>
+) : HashMap<Column<*>, Comparable<*>?>(idMapping), Comparable<CompositeID> {
+    override fun toString(): String = "CompositeID(${entries.joinToString { "${it.key.name}=${it.value}" }})"
+
+    override fun hashCode(): Int = entries.fold(0) { acc, entry ->
+        (acc * 31) + entry.hashCode()
+    }
+
+    override fun equals(other: Any?): Boolean {
+        if (other !is CompositeID) return false
+
+        return super.equals(other)
+    }
+
+    override fun compareTo(other: CompositeID): Int {
+        val compareSize = compareValues(other.size, size)
+        if (compareSize != 0) return compareSize
+
+        entries.forEach { (column, idValue) ->
+            if (!other.containsKey(column)) return -1
+            compareValues(idValue, other[column]).let {
+                if (it != 0) return it
+            }
+        }
+        return 0
+    }
+}

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

@@ -1,7 +1,5 @@
 package org.jetbrains.exposed.dao.id
 
-import org.jetbrains.exposed.sql.Column
-
 open class EntityID<T : Comparable<T>> protected constructor(val table: IdTable<T>, id: T?) : Comparable<EntityID<T>> {
     constructor(id: T, table: IdTable<T>) : this(table, id)
 
@@ -31,30 +29,3 @@ open class EntityID<T : Comparable<T>> protected constructor(val table: IdTable<
 
     override fun compareTo(other: EntityID<T>): Int = value.compareTo(other.value)
 }
-
-class CompositeID : HashMap<Column<*>, Comparable<*>?>(), Comparable<CompositeID> {
-    override fun toString(): String = "CompositeID(${entries.joinToString { "${it.key.name}=${it.value}" }})"
-
-    override fun hashCode(): Int = values.fold(0) { acc, id ->
-        (acc * 31) + id.hashCode()
-    }
-
-    override fun equals(other: Any?): Boolean {
-        if (other !is CompositeID) return false
-        val s = hashMapOf(1 to 3)
-
-        return super.equals(other)
-    }
-
-    override fun compareTo(other: CompositeID): Int {
-        val compareSize = compareValues(other.size, size)
-        if (compareSize != 0) return compareSize
-
-        entries.forEach { (column, id) ->
-            compareValues(id, other[column]).let {
-                if (it != 0) return it
-            }
-        }
-        return 0
-    }
-}

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

@@ -1,6 +1,7 @@
 package org.jetbrains.exposed.dao.id
 
 import org.jetbrains.exposed.sql.*
+import org.jetbrains.exposed.sql.SqlExpressionBuilder.wrap
 import java.util.*
 import kotlin.collections.HashSet
 
@@ -31,21 +32,50 @@ abstract class IdTable<T : Comparable<T>>(name: String = "") : Table(name) {
     abstract val id: Column<EntityID<T>>
 }
 
-abstract class CompositeIdTable(name: String = "") : IdTable<CompositeID>(name) {
+/**
+ * Identity table with a primary key consisting of a combination of columns.
+ *
+ * @param name Table name. By default, this will be resolved from any class name with a "Table" suffix removed (if present).
+ */
+open class CompositeIdTable(name: String = "") : IdTable<CompositeID>(name) {
+    /** The columns combined to make up this [CompositeIdTable]'s primary key. */
     val idColumns = HashSet<Column<out Comparable<*>>>()
 
-    fun compositeEntityId(
-        firstColumn: Column<out Comparable<*>>,
-        vararg columns: Column<out Comparable<*>>
-    ): Column<EntityID<CompositeID>> {
-        idColumns.addAll(arrayOf(firstColumn) + columns)
-        return Column<EntityID<CompositeID>>(this, "id", EntityIDColumnType(firstColumn as Column<out Comparable<Any>>)).also {
-            val compositeID = CompositeID().apply {
-                idColumns.forEach { column -> put(column, column.defaultValueFun?.let { it() }) }
+    final override val id: Column<EntityID<CompositeID>> = compositeIdColumn()
+
+    private fun compositeIdColumn(): Column<EntityID<CompositeID>> {
+        // Column class constructors are used to ensure neither column is actually registered in the DB via Table.columns
+        val placeholder = Column<String>(this, "composite_id", TextColumnType())
+        return Column<EntityID<CompositeID>>(this, "id", EntityIDColumnType(placeholder)).apply {
+            defaultValueFun = {
+                val defaultMap = idColumns.associateWith { column ->
+                    column.defaultValueFun?.let { it() }
+                }
+                EntityIDFunctionProvider.createEntityID(CompositeID(defaultMap), this@CompositeIdTable)
             }
-            it.defaultValueFun = { EntityIDFunctionProvider.createEntityID(compositeID, this) }
         }
     }
+
+    /** Marks [this] column as a component of a [CompositeIdTable]'s [EntityID] column. */
+    fun <T : Comparable<T>> Column<T>.compositeEntityId(): Column<T> = this.also { idColumns.add(it) }
+
+    /**
+     * Returns a list of boolean operators comparing each of this table's [idColumns] to its corresponding
+     * value in [toCompare], using the specified SQL [operator].
+     *
+     * @throws IllegalStateException If [toCompare] does not contain a key for each component column.
+     */
+    internal fun mapIdComparison(
+        toCompare: EntityID<CompositeID>,
+        operator: (Column<*>, Expression<*>) -> Op<Boolean>
+    ): List<Op<Boolean>> = idColumns.map { column ->
+        val otherValue = if (toCompare.value.containsKey(column)) {
+            toCompare.value[column]
+        } else {
+            error("Comparison CompositeID is missing a key mapping for ${column.name}")
+        }
+        operator(column, column.wrap(otherValue))
+    }
 }
 
 /**

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

@@ -27,18 +27,21 @@ class ResultRow(
     operator fun <T> get(expression: Expression<T>): T {
         val column = expression as? Column<*>
         return when {
-            column?.columnType is EntityIDColumnType<*> && column.table is CompositeIdTable -> getDeconstructedID(column.table.idColumns)
+            column?.columnType is EntityIDColumnType<*> && column.table is CompositeIdTable -> getIdComponents(column.table)
             else -> getInternal(expression, checkNullability = true)
         }
     }
 
-    private fun <T> getDeconstructedID(columns: Set<Column<*>>): T {
-        val compositeID = CompositeID().apply {
-            columns.forEach { column ->
-                put(column, getInternal(column, checkNullability = true) as Comparable<*>)
-            }
+    /**
+     * Retrieves the value for each component column from the specified [table] `id` column and returns the
+     * collective values as an [EntityID] value.
+     */
+    @Suppress("UNCHECKED_CAST")
+    private fun <T> getIdComponents(table: CompositeIdTable): T {
+        val resultMap = table.idColumns.associateWith { column ->
+            getInternal(column, checkNullability = true)
         }
-        return EntityID(compositeID, columns.first().table as CompositeIdTable) as T
+        return EntityID(CompositeID(resultMap), table) as T
     }
 
     /**

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

@@ -250,15 +250,12 @@ interface ISqlExpressionBuilder {
 
     /** Checks if this expression is equal to some [t] value. */
     @LowPriorityInOverloadResolution
+    @Suppress("UNCHECKED_CAST")
     infix fun <T> ExpressionWithColumnType<T>.eq(t: T): Op<Boolean> = when {
         t == null -> isNull()
         (columnType as? EntityIDColumnType<*>)?.idColumn?.table is CompositeIdTable -> {
             val table = (columnType as EntityIDColumnType<*>).idColumn.table as CompositeIdTable
-            val otherID = t as EntityID<CompositeID>
-            table.idColumns.map { column ->
-                val otherValue = if (otherID.value.containsKey(column)) otherID.value[column] else error("Comparison CompositeID missing value for ${column.name}")
-                (column as Column<Any?>) eq otherValue
-            }.compoundAnd()
+            table.mapIdComparison(t as EntityID<CompositeID>, ::EqOp).compoundAnd()
         }
         else -> EqOp(this, wrap(t))
     }
@@ -304,7 +301,15 @@ interface ISqlExpressionBuilder {
 
     /** Checks if this expression is not equal to some [other] value. */
     @LowPriorityInOverloadResolution
-    infix fun <T> ExpressionWithColumnType<T>.neq(other: T): Op<Boolean> = if (other == null) isNotNull() else NeqOp(this, wrap(other))
+    @Suppress("UNCHECKED_CAST")
+    infix fun <T> ExpressionWithColumnType<T>.neq(other: T): Op<Boolean> = when {
+        other == null -> isNotNull()
+        (columnType as? EntityIDColumnType<*>)?.idColumn?.table is CompositeIdTable -> {
+            val table = (columnType as EntityIDColumnType<*>).idColumn.table as CompositeIdTable
+            table.mapIdComparison(other as EntityID<CompositeID>, ::NeqOp).compoundAnd()
+        }
+        else -> NeqOp(this, wrap(other))
+    }
 
     /** Checks if this expression is not equal to some [other] expression. */
     infix fun <T, S1 : T?, S2 : T?> Expression<in S1>.neq(other: Expression<in S2>): Op<Boolean> = when (other as Expression<*>) {

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

@@ -2,6 +2,8 @@
 
 package org.jetbrains.exposed.sql
 
+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.EntityIDFunctionProvider
 import org.jetbrains.exposed.dao.id.IdTable
@@ -43,10 +45,15 @@ interface FieldSet {
             val unrolled = ArrayList<Expression<*>>(fields.size)
 
             fields.forEach {
-                if (it is CompositeColumn<*>) {
-                    unrolled.addAll(it.getRealColumns())
-                } else {
-                    unrolled.add(it)
+                when {
+                    it is CompositeColumn<*> -> unrolled.addAll(it.getRealColumns())
+                    (it as? Column<*>)?.columnType is EntityIDColumnType<*> -> {
+                        when (val table = (it as? Column<*>)?.table) {
+                            is CompositeIdTable -> unrolled.addAll(table.idColumns)
+                            else -> unrolled.add(it)
+                        }
+                    }
+                    else -> unrolled.add(it)
                 }
             }
 

exposed-dao/src/main/kotlin/org/jetbrains/exposed/dao/CompositeEntity.kt

@@ -4,9 +4,14 @@ import org.jetbrains.exposed.dao.id.CompositeID
 import org.jetbrains.exposed.dao.id.EntityID
 import org.jetbrains.exposed.dao.id.IdTable
 
+/** Base class for an [Entity] instance identified by an [id] comprised of multiple table column values. */
 abstract class CompositeEntity(id: EntityID<CompositeID>) : Entity<CompositeID>(id)
 
-abstract class CompositeEntityClass<out E : CompositeEntity> constructor(
+/**
+ * Base class representing the [EntityClass] that manages [CompositeEntity] instances and
+ * maintains their relation to the provided [table].
+ */
+abstract class CompositeEntityClass<out E : CompositeEntity>(
     table: IdTable<CompositeID>,
     entityType: Class<E>? = null,
     entityCtor: ((EntityID<CompositeID>) -> E)? = null

exposed-dao/src/main/kotlin/org/jetbrains/exposed/dao/Entity.kt

@@ -1,7 +1,10 @@
 package org.jetbrains.exposed.dao
 
 import org.jetbrains.exposed.dao.exceptions.EntityNotFoundException
+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.sql.*
 import org.jetbrains.exposed.sql.SqlExpressionBuilder.eq
 import org.jetbrains.exposed.sql.transactions.TransactionManager
@@ -303,4 +306,32 @@ open class Entity<ID : Comparable<ID>>(val id: EntityID<ID>) {
         // clear write values
         writeValues.clear()
     }
+
+    /**
+     * Stores a [value] for a [table] `id` column in this Entity's [writeValues] map.
+     * If the `id` column wraps a composite value, each non-null component value is stored for its component column.
+     */
+    @Suppress("UNCHECKED_CAST")
+    internal fun writeIdColumnValue(table: IdTable<*>, value: EntityID<*>) {
+        if (table is CompositeIdTable) {
+            val compositeID = value._value as CompositeID
+            table.idColumns.forEach { column ->
+                compositeID[column]?.let {
+                    writeValues[column as Column<Any?>] = it
+                }
+            }
+        } else {
+            writeValues[table.id as Column<Any?>] = value
+        }
+    }
+}
+
+/**
+ * Returns whether an [EntityID]'s wrapped value has not been assigned an actual value other than `null`.
+ * If [this] wraps a composite value, `true` is returned if any component is still assigned to `null`.
+ */
+internal fun EntityID<*>.valueIsNotInitialized(): Boolean = if (table is CompositeIdTable) {
+    (_value as CompositeID).values.any { it == null }
+} else {
+    _value == null
 }

exposed-dao/src/main/kotlin/org/jetbrains/exposed/dao/EntityCache.kt

@@ -194,15 +194,10 @@ class EntityCache(private val transaction: Transaction) {
             }
 
             for ((entry, genValues) in toFlush.zip(ids)) {
-                if (entry.id._value == null || (table is CompositeIdTable && (entry.id._value as CompositeID).values.any { it == null })) {
+                if (entry.id.valueIsNotInitialized()) {
                     val id = genValues[table.id]
                     entry.id._value = id._value
-                    when (val ekt = entry.klass.table) {
-                        is CompositeIdTable -> ekt.idColumns.forEach { column ->
-                            entry.writeValues[column as Column<Any?>] = id
-                        }
-                        else -> entry.writeValues[ekt.id as Column<Any?>] = id
-                    }
+                    entry.writeIdColumnValue(entry.klass.table, id)
                 }
                 genValues.fieldIndex.keys.forEach { key ->
                     entry.writeValues[key as Column<Any?>] = genValues[key]

exposed-dao/src/main/kotlin/org/jetbrains/exposed/dao/EntityClass.kt

@@ -301,23 +301,16 @@ abstract class EntityClass<ID : Comparable<ID>, out T : Entity<ID>>(
         prototype.klass = this
         prototype.db = TransactionManager.current().db
         prototype._readValues = ResultRow.createAndFillDefaults(dependsOnColumns)
-        if (entityId._value != null || table is CompositeIdTable) {
-            when (table) {
-                is CompositeIdTable -> (entityId._value as CompositeID).forEach { (column, idValue) ->
-                    idValue?.let {
-                        prototype.writeValues[column as Column<Any?>] = it
-                    }
-                }
-                else -> prototype.writeValues[table.id as Column<Any?>] = entityId
-            }
+        if (entityId._value != null) {
+            prototype.writeIdColumnValue(table, entityId)
         }
         try {
             entityCache.addNotInitializedEntityToQueue(prototype)
             prototype.init()
         } finally {
             entityCache.finishEntityInitialization(prototype)
         }
-        if (entityId._value == null || (table is CompositeIdTable && (entityId._value as CompositeID).values.any { it == null })) {
+        if (entityId.valueIsNotInitialized()) {
             val readValues = prototype._readValues!!
             val writeValues = prototype.writeValues
             table.columns.filter { col ->