Eager loading Parent-Child Reference #1363

JetBrains · Jul 25, 2022 · 267ff7d · 267ff7d
1 parent e926b0c
commit 267ff7d
Show file tree

Hide file tree

Showing 5 changed files with 116 additions and 51 deletions.
diff --git a/exposed-dao/src/main/kotlin/org/jetbrains/exposed/dao/Entity.kt b/exposed-dao/src/main/kotlin/org/jetbrains/exposed/dao/Entity.kt
@@ -228,13 +228,12 @@ open class Entity<ID : Comparable<ID>>(val id: EntityID<ID>) {
     }
 
     infix fun <TID : Comparable<TID>, Target : Entity<TID>> EntityClass<TID, Target>.via(table: Table): InnerTableLink<ID, Entity<ID>, TID, Target> =
-        InnerTableLink(table, this@via)
+        InnerTableLink(table, this@Entity.id.table, this@via)
 
     fun <TID : Comparable<TID>, Target : Entity<TID>> EntityClass<TID, Target>.via(
         sourceColumn: Column<EntityID<ID>>,
         targetColumn: Column<EntityID<TID>>
-    ) =
-        InnerTableLink(sourceColumn.table, this@via, sourceColumn, targetColumn)
+    ) = InnerTableLink(sourceColumn.table, this@Entity.id.table, this@via, sourceColumn, targetColumn)
 
     /**
      * Delete this entity.

diff --git a/exposed-dao/src/main/kotlin/org/jetbrains/exposed/dao/EntityClass.kt b/exposed-dao/src/main/kotlin/org/jetbrains/exposed/dao/EntityClass.kt
@@ -419,23 +419,16 @@ abstract class EntityClass<ID : Comparable<ID>, out T : Entity<ID>>(
         else -> findQuery
     }.toList()
 
-    /**
-     * @param optimizedLoad will force to make to two queries to load ids and referenced entities separately.
-     * Can be useful when references target the same entities. That will prevent from loading them multiple times (per each reference row) and will require
-     * less memory/bandwidth for "heavy" entities (with a lot of columns or columns with huge data in it)
-     */
-    fun warmUpLinkedReferences(references: List<EntityID<*>>, linkTable: Table, forUpdate: Boolean? = null, optimizedLoad: Boolean = false): List<T> {
+    internal fun <SID: Comparable<SID>> warmUpLinkedReferences(
+        references: List<EntityID<SID>>, sourceRefColumn: Column<EntityID<SID>>, targetRefColumn: Column<EntityID<ID>>, linkTable: Table,
+        forUpdate: Boolean? = null, optimizedLoad: Boolean = false
+    ) : List<T> {
         if (references.isEmpty()) return emptyList()
         val distinctRefIds = references.distinct()
-        val sourceRefColumn = linkTable.columns.singleOrNull { it.referee == references.first().table.id } as? Column<EntityID<*>>
-            ?: error("Can't detect source reference column")
-        val targetRefColumn =
-            linkTable.columns.singleOrNull { it.referee == table.id } as? Column<EntityID<ID>> ?: error("Can't detect target reference column")
-
         val transaction = TransactionManager.current()
 
         val inCache = transaction.entityCache.referrers[sourceRefColumn] ?: emptyMap()
-        val loaded = (distinctRefIds - inCache.keys).takeIf { it.isNotEmpty() }?.let { idsToLoad ->
+        val loaded = ((distinctRefIds - inCache.keys).takeIf { it.isNotEmpty() } as List<EntityID<SID>>?)?.let { idsToLoad ->
             val alreadyInJoin = (dependsOnTables as? Join)?.alreadyInJoin(linkTable) ?: false
             val entityTables = if (alreadyInJoin) dependsOnTables else dependsOnTables.join(linkTable, JoinType.INNER, targetRefColumn, table.id)
 
@@ -453,7 +446,7 @@ abstract class EntityClass<ID : Comparable<ID>, out T : Entity<ID>>(
                 else -> query
             }.map {
                 val targetId = it[targetRefColumn]
-                 if (!optimizedLoad) {
+                if (!optimizedLoad) {
                     targetEntities.getOrPut(targetId) { wrapRow(it) }
                 }
                 it[sourceRefColumn] to targetId
@@ -477,9 +470,27 @@ abstract class EntityClass<ID : Comparable<ID>, out T : Entity<ID>>(
         return inCache.values.flatMap { it.toList() as List<T> } + loaded.orEmpty()
     }
 
+    /**
+     * @param optimizedLoad will force to make to two queries to load ids and referenced entities separately.
+     * Can be useful when references target the same entities. That will prevent from loading them multiple times (per each reference row) and will require
+     * less memory/bandwidth for "heavy" entities (with a lot of columns or columns with huge data in it)
+     */
+    fun <SID: Comparable<SID>> warmUpLinkedReferences(references: List<EntityID<SID>>, linkTable: Table, forUpdate: Boolean? = null, optimizedLoad: Boolean = false): List<T> {
+        if (references.isEmpty()) return emptyList()
+
+        val sourceRefColumn = linkTable.columns.singleOrNull { it.referee == references.first().table.id } as? Column<EntityID<SID>>
+            ?: error("Can't detect source reference column")
+        val targetRefColumn =
+            linkTable.columns.singleOrNull { it.referee == table.id } as? Column<EntityID<ID>> ?: error("Can't detect target reference column")
+
+        return warmUpLinkedReferences(references, sourceRefColumn, targetRefColumn, linkTable, forUpdate, optimizedLoad)
+    }
+
     fun <ID : Comparable<ID>, T : Entity<ID>> isAssignableTo(entityClass: EntityClass<ID, T>) = entityClass.klass.isAssignableFrom(klass)
 }
 
+
+
 abstract class ImmutableEntityClass<ID : Comparable<ID>, out T : Entity<ID>>(
     table: IdTable<ID>,
     entityType: Class<T>? = null,

diff --git a/exposed-dao/src/main/kotlin/org/jetbrains/exposed/dao/InnerTableLink.kt b/exposed-dao/src/main/kotlin/org/jetbrains/exposed/dao/InnerTableLink.kt
@@ -1,6 +1,7 @@
 package org.jetbrains.exposed.dao
 
 import org.jetbrains.exposed.dao.id.EntityID
+import org.jetbrains.exposed.dao.id.IdTable
 import org.jetbrains.exposed.sql.*
 import org.jetbrains.exposed.sql.transactions.TransactionManager
 import kotlin.properties.ReadWriteProperty
@@ -9,52 +10,57 @@ import kotlin.reflect.KProperty
 @Suppress("UNCHECKED_CAST")
 class InnerTableLink<SID : Comparable<SID>, Source : Entity<SID>, ID : Comparable<ID>, Target : Entity<ID>>(
     val table: Table,
+    sourceTable: IdTable<SID>,
     val target: EntityClass<ID, Target>,
-    val sourceColumn: Column<EntityID<SID>>? = null,
-    targetColumn: Column<EntityID<ID>>? = null
+    _sourceColumn: Column<EntityID<SID>>? = null,
+    _targetColumn: Column<EntityID<ID>>? = null,
 ) : ReadWriteProperty<Source, SizedIterable<Target>> {
     init {
-        targetColumn?.let {
-            requireNotNull(sourceColumn) { "Both source and target columns should be specified" }
-            require(targetColumn.referee?.table == target.table) {
-                "Column $targetColumn point to wrong table, expected ${target.table.tableName}"
+        _targetColumn?.let {
+            requireNotNull(_sourceColumn) { "Both source and target columns should be specified" }
+            require(_targetColumn.referee?.table == target.table) {
+                "Column $_targetColumn point to wrong table, expected ${target.table.tableName}"
             }
-            require(targetColumn.table == sourceColumn.table) {
+            require(_targetColumn.table == _sourceColumn.table) {
                 "Both source and target columns should be from the same table"
             }
         }
-        sourceColumn?.let {
-            requireNotNull(targetColumn) { "Both source and target columns should be specified" }
+        _sourceColumn?.let {
+            requireNotNull(_targetColumn) { "Both source and target columns should be specified" }
+            require(_sourceColumn.referee?.table == sourceTable) {
+                "Column $_sourceColumn point to wrong table, expected ${sourceTable.tableName}"
+            }
         }
     }
 
-    private val targetColumn = targetColumn
+    val sourceColumn = _sourceColumn
+        ?: table.columns.singleOrNull { it.referee == sourceTable.id } as? Column<EntityID<SID>>
+        ?: error("Table does not reference source")
+
+    val targetColumn = _targetColumn
         ?: table.columns.singleOrNull { it.referee == target.table.id } as? Column<EntityID<ID>>
         ?: error("Table does not reference target")
 
-    private fun getSourceRefColumn(o: Source): Column<EntityID<SID>> {
-        return sourceColumn
-            ?: table.columns.singleOrNull { it.referee == o.klass.table.id } as? Column<EntityID<SID>>
-            ?: error("Table does not reference source")
+    private val columnsAndTables by lazy {
+        val alreadyInJoin = (target.dependsOnTables as? Join)?.alreadyInJoin(table) ?: false
+        val entityTables =
+            if (alreadyInJoin) target.dependsOnTables else target.dependsOnTables.join(table, JoinType.INNER, target.table.id, targetColumn)
+
+        val columns = (target.dependsOnColumns + (if (!alreadyInJoin) table.columns else emptyList()) - sourceColumn).distinct() + sourceColumn
+
+        columns to entityTables
     }
 
     override operator fun getValue(o: Source, unused: KProperty<*>): SizedIterable<Target> {
         if (o.id._value == null && !o.isNewEntity()) return emptySized()
-        val sourceRefColumn = getSourceRefColumn(o)
         val transaction = TransactionManager.currentOrNull()
-            ?: return o.getReferenceFromCache(sourceRefColumn)
-        val alreadyInJoin = (target.dependsOnTables as? Join)?.alreadyInJoin(table) ?: false
-        val entityTables =
-            if (alreadyInJoin) target.dependsOnTables else target.dependsOnTables.join(table, JoinType.INNER, target.table.id, targetColumn)
+            ?: return o.getReferenceFromCache(sourceColumn)
 
-        val columns = (
-            target.dependsOnColumns + (if (!alreadyInJoin) table.columns else emptyList()) -
-                sourceRefColumn
-            ).distinct() + sourceRefColumn
+        val (columns, entityTables) = columnsAndTables
 
-        val query = { target.wrapRows(entityTables.slice(columns).select { sourceRefColumn eq o.id }) }
-        return transaction.entityCache.getOrPutReferrers(o.id, sourceRefColumn, query).also {
-            o.storeReferenceInCache(sourceRefColumn, it)
+        val query = { target.wrapRows(entityTables.slice(columns).select { sourceColumn eq o.id }) }
+        return transaction.entityCache.getOrPutReferrers(o.id, sourceColumn, query).also {
+            o.storeReferenceInCache(sourceColumn, it)
         }
     }
 
@@ -70,20 +76,18 @@ class InnerTableLink<SID : Comparable<SID>, Source : Entity<SID>, ID : Comparabl
     }
 
     private fun setReference(o: Source, unused: KProperty<*>, value: SizedIterable<Target>) {
-        val sourceRefColumn = getSourceRefColumn(o)
-
         val tx = TransactionManager.current()
         val entityCache = tx.entityCache
         entityCache.flush()
         val oldValue = getValue(o, unused)
         val existingIds = oldValue.map { it.id }.toSet()
-        entityCache.referrers[sourceRefColumn]?.remove(o.id)
+        entityCache.referrers[sourceColumn]?.remove(o.id)
 
         val targetIds = value.map { it.id }
         executeAsPartOfEntityLifecycle {
-            table.deleteWhere { (sourceRefColumn eq o.id) and (targetColumn notInList targetIds) }
+            table.deleteWhere { (sourceColumn eq o.id) and (targetColumn notInList targetIds) }
             table.batchInsert(targetIds.filter { !existingIds.contains(it) }, shouldReturnGeneratedValues = false) { targetId ->
-                this[sourceRefColumn] = o.id
+                this[sourceColumn] = o.id
                 this[targetColumn] = targetId
             }
         }

diff --git a/exposed-dao/src/main/kotlin/org/jetbrains/exposed/dao/References.kt b/exposed-dao/src/main/kotlin/org/jetbrains/exposed/dao/References.kt
@@ -178,9 +178,15 @@ private fun <ID : Comparable<ID>> List<Entity<ID>>.preloadRelations(
                 }
             }
             is InnerTableLink<*, *, *, *> -> {
-                refObject.target.warmUpLinkedReferences(this.map { it.id }, refObject.table)
-                val refColumn = refObject.table.columns.single { it.referee == this.first().id.table.id }
-                storeReferenceCache(refColumn, prop)
+                (refObject as InnerTableLink<ID, Entity<ID>, Comparable<Comparable<*>>, Entity<Comparable<Comparable<*>>>>).let { innerTableLink ->
+                    innerTableLink.target.warmUpLinkedReferences(
+                        references = this.map { it.id },
+                        sourceRefColumn = innerTableLink.sourceColumn,
+                        targetRefColumn = innerTableLink.targetColumn,
+                        linkTable = innerTableLink.table
+                    )
+                    storeReferenceCache(innerTableLink.sourceColumn, prop)
+                }
             }
             is BackReference<*, *, *, *, *> -> {
                 (refObject.delegate as Referrers<ID, Entity<ID>, *, Entity<*>, Any>).reference.let { refColumn ->

diff --git a/exposed-tests/src/test/kotlin/org/jetbrains/exposed/sql/tests/shared/entities/ViaTest.kt b/exposed-tests/src/test/kotlin/org/jetbrains/exposed/sql/tests/shared/entities/ViaTest.kt
@@ -8,9 +8,11 @@ import org.jetbrains.exposed.dao.id.UUIDTable
 import org.jetbrains.exposed.sql.*
 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.junit.Test
 import java.util.*
+import kotlin.reflect.jvm.isAccessible
 
 object ViaTestData {
     object NumbersTable : UUIDTable() {
@@ -165,12 +167,13 @@ class ViaTests : DatabaseTestsBase() {
         var name by NodesTable.name
         var parents by Node.via(NodeToNodes.child, NodeToNodes.parent)
         var children by Node.via(NodeToNodes.parent, NodeToNodes.child)
+
+        override fun equals(other: Any?): Boolean = (other as? Node)?.id == id
     }
 
     @Test
     fun testHierarchicalReferences() {
         withTables(NodeToNodes) {
-            addLogger(StdOutSqlLogger)
             val root = Node.new { name = "root" }
             val child1 = Node.new {
                 name = "child1"
@@ -196,4 +199,46 @@ class ViaTests : DatabaseTestsBase() {
             assertEquals("ccc", s.text)
         }
     }
+
+    @Test
+    fun testWarmUpOnHierarchicalEntities() {
+        withTables(NodeToNodes) {
+            val child1 = Node.new { name = "child1" }
+            val child2 = Node.new { name = "child1" }
+            val root1 = Node.new {
+                name = "root1"
+                children = SizedCollection(child1)
+            }
+            val root2 = Node.new {
+                name = "root2"
+                children = SizedCollection(child1, child2)
+            }
+
+            entityCache.clear(flush = true)
+
+            fun checkChildrenReferences(node: Node, values: List<Node>) {
+                val sourceColumn = (Node::children.apply { isAccessible = true }.getDelegate(node) as InnerTableLink<*,*,*,*>).sourceColumn
+                val children = entityCache.getReferrers<Node>(node.id, sourceColumn)
+                assertEqualLists(children?.toList().orEmpty(), values)
+            }
+
+            Node.all().with(Node::children).toList()
+            checkChildrenReferences(child1, emptyList())
+            checkChildrenReferences(child2, emptyList())
+            checkChildrenReferences(root1, listOf(child1))
+            checkChildrenReferences(root2, listOf(child1, child2))
+
+            fun checkParentsReferences(node: Node, values: List<Node>) {
+                val sourceColumn =  (Node::parents.apply { isAccessible = true }.getDelegate(node) as InnerTableLink<*,*,*,*>).sourceColumn
+                val children = entityCache.getReferrers<Node>(node.id, sourceColumn)
+                assertEqualLists(children?.toList().orEmpty(), values)
+            }
+
+            Node.all().with(Node::parents).toList()
+            checkParentsReferences(child1, listOf(root1, root2))
+            checkParentsReferences(child2, listOf(root2))
+            checkParentsReferences(root1, emptyList())
+            checkParentsReferences(root2, emptyList())
+        }
+    }
 }