add parameter for type replacements (decouple driver Time)

2 files changed, 15 insertions, 12 deletions

diff --git a/src/main/scala/CodegenPlugin.scala b/src/main/scala/CodegenPlugin.scala
index 1286919..b3a97ca 100644
--- a/src/main/scala/CodegenPlugin.scala
+++ b/src/main/scala/CodegenPlugin.scala
@@ -32,6 +32,10 @@ object CodegenPlugin extends AutoPlugin {
       "codegen-schema-imports",
       "A list of things to import into each schema definition"
     )
+    lazy val codegenTypeReplacements = SettingKey[Map[String, String]](
+      "codegen-type-replacements",
+      "A map of types to find and replace"
+    )
 
     lazy val slickCodeGenTask =
       TaskKey[Unit]("gen-tables", "generate the table definitions")
@@ -46,6 +50,7 @@ object CodegenPlugin extends AutoPlugin {
     codegenSchemaBaseClassParts := List.empty,
     codegenIdType := Option.empty,
     codegenSchemaImports := List.empty,
+    codegenTypeReplacements := Map.empty,
     slickCodeGenTask := Def.taskDyn {
       Def.task {
         Generator.run(
@@ -62,7 +67,8 @@ object CodegenPlugin extends AutoPlugin {
             case parts => parts.mkString(" with ")
           },
           codegenIdType.value,
-          codegenSchemaImports.value
+          codegenSchemaImports.value,
+          codegenTypeReplacements.value
         )
       }
     }.value
diff --git a/src/main/scala/NamespacedCodegen.scala b/src/main/scala/NamespacedCodegen.scala
index fb4910a..a3df345 100644
--- a/src/main/scala/NamespacedCodegen.scala
+++ b/src/main/scala/NamespacedCodegen.scala
@@ -25,7 +25,8 @@ object Generator {
           manualForeignKeys: Map[(String, String), (String, String)],
           schemaBaseClass: String,
           idType: Option[String],
-          schemaImports: List[String]) = {
+    schemaImports: List[String],
+  typeReplacements: Map[String, String]) = {
     val dc: DatabaseConfig[JdbcProfile] =
       DatabaseConfig.forURI[JdbcProfile](uri)
     val parsedSchemasOpt: Option[Map[String, List[String]]] =
@@ -41,7 +42,8 @@ object Generator {
                                   manualForeignKeys,
                                   schemaBaseClass,
                                   idType,
-                                  schemaImports)
+      schemaImports,
+    typeReplacements)
     generator.code // Yes... Files are written as a side effect
     parsedSchemasOpt
       .getOrElse(Map())
@@ -96,7 +98,8 @@ class Generator(uri: URI,
                 manualForeignKeys: Map[(String, String), (String, String)],
                 schemaBaseClass: String,
                 idType: Option[String],
-                schemaImports: List[String])
+  schemaImports: List[String],
+  typeReplacements: Map[String, String])
     extends SourceCodeGenerator(dbModel)
     with OutputHelpers {
 
@@ -226,20 +229,14 @@ class Generator(uri: URI,
         s"$idTypeName[$schemaObjectName.$rowTypeName]"
       }
 
-      // re-write ids, and time types
+      // re-write ids other custom types
       override def rawType: String = {
         val (referencedTable, referencedColumn) =
           derefColumn(table.model, column.model)
         if (referencedColumn.options.contains(
               slick.ast.ColumnOption.PrimaryKey))
           tableReferenceName(referencedTable.name)
-        else
-          model.tpe match {
-            // TODO: There should be a way to add adhoc custom time mappings
-            case "java.sql.Time" => "xyz.driver.core.time.Time"
-            case "java.sql.Timestamp" => "xyz.driver.core.time.Time"
-            case _ => super.rawType
-          }
+        else typeReplacements.getOrElse(model.tpe, model.tpe)
       }
     }