add primary key mapper TermDef

2 files changed, 21 insertions, 17 deletions

diff --git a/src/main/scala/Generators.scala b/src/main/scala/Generators.scala
index 9af3d88..7138a3c 100644
--- a/src/main/scala/Generators.scala
+++ b/src/main/scala/Generators.scala
@@ -20,8 +20,6 @@ class RowSourceCodeGenerator(
 
   override def Table = new TypedIdTable(_) { table =>
 
-    override def PrimaryKey = new TypedIdPrimaryKey(_) { }
-
     override def Column = new TypedIdColumn(_) {
       override def rawType: String = {
         typeReplacements.getOrElse(model.tpe, super.rawType)
diff --git a/src/main/scala/TypedIdTable.scala b/src/main/scala/TypedIdTable.scala
index 0958ccf..81a6dcf 100644
--- a/src/main/scala/TypedIdTable.scala
+++ b/src/main/scala/TypedIdTable.scala
@@ -42,27 +42,33 @@ class TypedIdSourceCodeGenerator(
     s"${idType.getOrElse("Id")}[${pKeyTypeTag(columnRef)}]"
   }
 
-  class TypedIdTable(model: m.Table) extends Table(model) {
+  class TypedIdTable(model: m.Table) extends Table(model) { table =>
+    override def definitions = Seq[Def]( EntityType, PlainSqlMapper, TableClass, TableValue, PrimaryKeyMapper )
+
     class TypedIdColumn(override val model: m.Column) extends Column(model) {
       override def rawType: String = {
         keyReferences.get(model).fold(super.rawType)(pKeyType)
       }
     }
 
-    class TypedIdPrimaryKey(override val model: m.PrimaryKey) extends PrimaryKey(model) { primaryKey =>
-      def `super.code` = s"""val $name = primaryKey("$dbName", ${compoundValue(columns.map(_.name))})"""
-
-      override def code = {
-        val implicitKeyBaseMapper =
-        primaryKey.columns.headOption
-          .filter(_ => primaryKey.columns.length == 1)
-          .map { column =>
-            val name = termName(column.rawName + "KeyMapper")
-            val tpe = s"BaseColumnType[column.rawName]"
-            val mapping = s"${modelTypeToColumnMaper(column.model.tpe)}[${pKeyTypeTag(column.model)}]"
-            s"implicit def $name: $tpe = $mapping\n"
-        }
-        implicitKeyBaseMapper.fold(super.code)(super.code + _)
+    type PrimaryKeyMapper = PrimaryKeyMapperDef
+
+    def PrimaryKeyMapper = new PrimaryKeyMapper { }
+
+    class PrimaryKeyMapperDef extends TermDef {
+      def primaryKeyColumn: Option[Column] = table.primaryKey.filter(_.columns.length == 1).flatMap(_.columns.headOption)
+
+      override def enabled = primaryKeyColumn.isDefined
+
+      override def doc = s"Implicit for mapping primary key of ${tableName} to a base column"
+
+      override def rawName = tableName+"KeyMapper"
+
+      override def code = primaryKeyColumn.fold("") { column =>
+        val tpe = s"BaseColumnType[${column.rawName}]"
+        s"""|implicit def $name: $tpe = 
+            |${modelTypeToColumnMaper(column.model.tpe)}[${pKeyTypeTag(column.model)}]
+            |""".stripMargin.lines.mkString("").trim
       }
     }
   }