Merge pull request #29 from drivergroup/row-import-statements

Add parameter for row-file specific imports

2 files changed, 11 insertions, 4 deletions

diff --git a/src/main/scala/CodegenPlugin.scala b/src/main/scala/CodegenPlugin.scala
index c90fe66..1ede3b8 100644
--- a/src/main/scala/CodegenPlugin.scala
+++ b/src/main/scala/CodegenPlugin.scala
@@ -42,6 +42,10 @@ object CodegenPlugin extends AutoPlugin {
       "codegen-schema-imports",
       "A list of things to import into each schema definition"
     )
+    lazy val codegenSchemaRowImports = SettingKey[Option[List[String]]](
+      "codegen-schema-imports",
+      "An optional list of things to import for table row definitions of each schema. Uses schemaImports otherwise."
+    )
     lazy val codegenTypeReplacements = SettingKey[Map[String, String]](
       "codegen-type-replacements",
       "A map of types to find and replace"
@@ -84,6 +88,7 @@ object CodegenPlugin extends AutoPlugin {
               codegenIdType.value,
               codegenHeader.value,
               codegenSchemaImports.value,
+              codegenSchemaRowImports.value getOrElse codegenSchemaImports.value,
               codegenTypeReplacements.value
             )
         }
diff --git a/src/main/scala/Main.scala b/src/main/scala/Main.scala
index 2602103..e71a9e8 100644
--- a/src/main/scala/Main.scala
+++ b/src/main/scala/Main.scala
@@ -46,13 +46,15 @@ object Generator {
           parentType: Option[String],
           idType: Option[String],
           header: String,
-          schemaImports: List[String],
+          tablesFileImports: List[String],
+          rowsFileImports: List[String],
           typeReplacements: Map[String, String]) = {
     val dc: DatabaseConfig[JdbcProfile] =
       DatabaseConfig.forURI[JdbcProfile](uri)
     val parsedSchemasOpt: Option[Map[String, List[String]]] =
       schemaNames.map(ModelTransformation.parseSchemaList)
-    val imports = schemaImports.map("import " + _).mkString("\n")
+
+    def importStatements(imports: List[String]) = imports.map("import " + _).mkString("\n")
 
     try {
       val dbModel: slick.model.Model = Await.result(
@@ -74,7 +76,7 @@ object Generator {
           val rowGenerator = new RowSourceCodeGenerator(
             model = schemaOnlyModel,
             headerComment = header,
-            imports = imports,
+            imports = importStatements(rowsFileImports),
             schemaName = schemaName,
             fullDatabaseModel = dbModel,
             idType,
@@ -85,7 +87,7 @@ object Generator {
           val tableGenerator =
             new TableSourceCodeGenerator(schemaOnlyModel = schemaOnlyModel,
                                          headerComment = header,
-                                         imports = imports,
+                                         imports = importStatements(tablesFileImports),
                                          schemaName = schemaName,
                                          fullDatabaseModel = dbModel,
                                          pkg = pkg,