From faccc16b0db9cd132a8e31fe58fee18d4f131de1 Mon Sep 17 00:00:00 2001
From: Stewart Stewart <stewinsalot@gmail.com>
Date: Tue, 22 Nov 2016 20:05:39 -0500
Subject: add paramater for list of imports per schema

---
 src/main/scala/CodegenPlugin.scala     |  8 +++++++-
 src/main/scala/NamespacedCodegen.scala | 21 +++++++++++----------
 2 files changed, 18 insertions(+), 11 deletions(-)

(limited to 'src/main')

diff --git a/src/main/scala/CodegenPlugin.scala b/src/main/scala/CodegenPlugin.scala
index 8b54887..01286ee 100644
--- a/src/main/scala/CodegenPlugin.scala
+++ b/src/main/scala/CodegenPlugin.scala
@@ -28,6 +28,10 @@ object CodegenPlugin extends AutoPlugin {
       "codegen-id-type",
       "The in-scope type `T` of kind `T[TableRow]` to apply in place T for id columns"
     )
+    lazy val codegenSchemaImports = SettingKey[List[String]](
+      "codegen-schema-imports",
+      "A list of things to import into each schema definition"
+    )
 
     lazy val slickCodeGenTask =
       TaskKey[Unit]("gen-tables", "generate the table definitions")
@@ -41,6 +45,7 @@ object CodegenPlugin extends AutoPlugin {
     codegenForeignKeys := Map.empty,
     codegenSchemaBaseClassParts := List.empty,
     codegenIdType := Option.empty,
+    codegenSchemaImports := List.empty,
     slickCodeGenTask := Def.taskDyn {
       Def.task {
         Generator.run(
@@ -53,7 +58,8 @@ object CodegenPlugin extends AutoPlugin {
             case Nil => "AnyRef"
             case parts => parts.mkString(" with ")
           },
-          codegenIdType.value
+          codegenIdType.value,
+          codegenSchemaImports.value
         )
       }
     }.value
diff --git a/src/main/scala/NamespacedCodegen.scala b/src/main/scala/NamespacedCodegen.scala
index bddbdea..2ac41e0 100644
--- a/src/main/scala/NamespacedCodegen.scala
+++ b/src/main/scala/NamespacedCodegen.scala
@@ -24,7 +24,8 @@ object Generator {
           outputPath: String,
           manualForeignKeys: Map[(String, String), (String, String)],
           schemaBaseClass: String,
-          idType: Option[String]) = {
+          idType: Option[String],
+          schemaImports: List[String]) = {
     val dc: DatabaseConfig[JdbcProfile] =
       DatabaseConfig.forURI[JdbcProfile](uri)
     val parsedSchemasOpt: Option[Map[String, List[String]]] =
@@ -39,7 +40,8 @@ object Generator {
                                   outputPath,
                                   manualForeignKeys,
                                   schemaBaseClass,
-                                  idType)
+                                  idType,
+                                  schemaImports)
     generator.code // Yes... Files are written as a side effect
     parsedSchemasOpt
       .getOrElse(Map())
@@ -60,12 +62,10 @@ class PackageNameGenerator(pkg: String, dbModel: Model)
        |""".stripMargin
 }
 
-class ImportGenerator(dbModel: Model) extends SourceCodeGenerator(dbModel) {
-  val baseImports: String =
-    s"""
-       |import xyz.driver.core.Id
-       |
-    |""".stripMargin
+class ImportGenerator(dbModel: Model, schemaImports: List[String])
+    extends SourceCodeGenerator(dbModel) {
+
+  val baseImports: String = schemaImports.map("import " + _).mkString("\n")
 
   val hlistImports: String =
     if (tables.exists(_.hlistEnabled))
@@ -95,12 +95,13 @@ class Generator(uri: URI,
                 outputPath: String,
                 manualForeignKeys: Map[(String, String), (String, String)],
                 schemaBaseClass: String,
-                idType: Option[String])
+                idType: Option[String],
+                schemaImports: List[String])
     extends SourceCodeGenerator(dbModel)
     with OutputHelpers {
 
   val packageName = new PackageNameGenerator(pkg, dbModel).code
-  val allImports: String = new ImportGenerator(dbModel).code
+  val allImports: String = new ImportGenerator(dbModel, schemaImports).code
 
   private val defaultIdImplementation =
     """|case class Id[T](v: Int)
-- 
cgit v1.2.3