pre-calculate all key references

1 files changed, 18 insertions, 19 deletions

diff --git a/src/main/scala/TypedIdTable.scala b/src/main/scala/TypedIdTable.scala
index 93ea8ee..740158d 100644
--- a/src/main/scala/TypedIdTable.scala
+++ b/src/main/scala/TypedIdTable.scala
@@ -16,35 +16,34 @@ class TypedIdSourceCodeGenerator(
     "Int" -> "serialKeyMapper"
   )
 
+  val keyReferences: Map[m.Column, m.Column] = {
+    val pks = databaseModel.tables.flatMap(_.columns)
+      .filter(_.options.contains(slick.ast.ColumnOption.PrimaryKey))
+      .map(c => (c -> c))
+
+    val fks: Seq[(m.Column, m.Column)] = databaseModel.tables.flatMap(_.foreignKeys)
+      .filter(_.referencedColumns.length == 1)
+      .filter(_.referencedColumns.forall(
+        _.options.contains(slick.ast.ColumnOption.PrimaryKey)))
+      .flatMap(fk =>
+      fk.referencingColumns.flatMap(from =>
+        fk.referencedColumns.headOption.map(to => (from -> to))))
+
+    (pks ++ fks).toMap
+  }
+
   def pKeyTypeTag(columnRef: m.Column): String = {
     val schemaName = columnRef.table.schema.getOrElse("`public`")
     val tableName = entityName(columnRef.table.table)
-    s"$schemaName.$tableName."
+    s"$schemaName.$tableName"
   }
 
   def pKeyType(columnRef: m.Column): String = {
     s"${idType.getOrElse("Id")}[${pKeyTypeTag(columnRef)}]"
   }
 
-  class TypedIdTable(model: m.Table) extends Table(model) { table =>
-
-    val keyReferences: Map[m.Column, m.Column] = {
-      val fks = model.foreignKeys
-        .filter(_.referencedColumns.length == 1)
-        .filter(_.referencedColumns.forall(
-          _.options.contains(slick.ast.ColumnOption.PrimaryKey)))
-        .flatMap(fk =>
-          fk.referencingColumns.flatMap(from =>
-            fk.referencedColumns.headOption.map(to => (from -> to))))
-      val pk = model.primaryKey
-        .filter(_.columns.length == 1)
-        .flatMap(_.columns.headOption.map(c => (c -> c)))
-
-      fks.toMap ++ pk
-    }
-
+  class TypedIdTable(model: m.Table) extends Table(model) {
     class TypedIdColumn(override val model: m.Column) extends Column(model) {
-      column =>
       override def rawType: String = {
         keyReferences.get(model).fold(super.rawType)(pKeyType)
       }