support for joinTable when columns are from the same source

3ricL · Apr 30, 2017 · fff5f6f · fff5f6f
1 parent 8a26717
commit fff5f6f
Showing 1 changed file with 9 additions and 3 deletions.
diff --git a/src/main/kotlin/org/jetbrains/exposed/dao/Entity.kt b/src/main/kotlin/org/jetbrains/exposed/dao/Entity.kt
@@ -117,14 +117,16 @@ class View<out Target: Entity<*>> (val op : Op<Boolean>, val factory: EntityClas
 
 @Suppress("UNCHECKED_CAST")
 class InnerTableLink<ID:Any, Target: Entity<ID>>(val table: Table,
-                                     val target: EntityClass<ID, Target>) {
+                                                 val target: EntityClass<ID, Target>,
+                                                 val source: Column<*>? = null) {
+
     private fun getSourceRefColumn(o: Entity<*>): Column<EntityID<*>> {
-        val sourceRefColumn = table.columns.singleOrNull { it.referee == o.klass.table.id } as? Column<EntityID<*>> ?: error("Table does not reference source")
+        val sourceRefColumn = table.columns.singleOrNull { it.referee == o.klass.table.id && if(source == null) true else it == source } as? Column<EntityID<*>> ?: error("Table does not reference source")
         return sourceRefColumn
     }
 
     private fun getTargetRefColumn(): Column<EntityID<*>> {
-        val sourceRefColumn = table.columns.singleOrNull { it.referee == target.table.id } as? Column<EntityID<*>> ?: error("Table does not reference target")
+        val sourceRefColumn = table.columns.singleOrNull { it.referee == target.table.id && if(source == null) true else it != source } as? Column<EntityID<*>> ?: error("Table does not reference target")
         return sourceRefColumn
     }
 
@@ -251,6 +253,10 @@ open class Entity<ID:Any>(val id: EntityID<ID>) {
         return InnerTableLink(table, this@via)
     }
 
+    infix fun <ID:Any, Target:Entity<ID>> EntityClass<ID, Target>.via(sourceColumn: Column<EntityID<ID>>): InnerTableLink<ID, Target> {
+        return InnerTableLink(sourceColumn.table, this@via, sourceColumn)
+    }
+
     fun <T: Entity<*>> s(c: EntityClass<T, *>): EntityClass<T, *> = c
 
     open fun delete(){