codegen tests

build.sbt

@@ -34,6 +34,9 @@ lazy val `slick-additions-codegen` =
     .settings(
       libraryDependencies ++= Seq(
         "com.typesafe.slick" %% "slick-hikaricp" % slickVersion,
-        "org.scalameta"      %% "scalameta"      % "4.8.14"
+        "org.scalameta"      %% "scalameta"      % "4.8.14",
+        "org.scalameta"      %% "scalafmt-core"  % "3.7.17",
+        "com.h2database"      % "h2"             % "2.2.224" % "test",
+        "org.scalatest"      %% "scalatest"      % "3.2.17"  % "test"
       )
     )

slick-additions-codegen/src/main/scala/slick/additions/codegen/BaseCodeGenerator.scala

@@ -3,13 +3,15 @@ package slick.additions.codegen
 import java.nio.file.{Files, Path}
 
 import scala.concurrent.duration.Duration
-import scala.concurrent.{Await, ExecutionContext}
+import scala.concurrent.{Await, ExecutionContext, Future}
 import scala.meta._
 
 import slick.dbio.DBIO
 import slick.jdbc.{JdbcBackend, JdbcProfile}
 
 import com.typesafe.config.Config
+import org.scalafmt.Scalafmt
+import org.scalafmt.config.ScalafmtConfig
 
 
 /** Base trait for code generators. Code generators are responsible for producing actual code, but many of the details
@@ -54,18 +56,43 @@ trait BaseCodeGenerator {
   def codeString(rules: GenerationRules, slickProfileClassName: String)(implicit executionContext: ExecutionContext)
     : DBIO[String] = codeString(rules, Class.forName(slickProfileClassName).asSubclass(classOf[JdbcProfile]))
 
-  def doWriteToFile(
+  def codeStringFormatted(
+    rules: GenerationRules,
+    slickProfileClassName: String,
+    scalafmtConfig: ScalafmtConfig = ScalafmtConfig.defaultWithAlign
+  )(implicit executionContext: ExecutionContext
+  ): DBIO[String] =
+    codeString(rules, Class.forName(slickProfileClassName).asSubclass(classOf[JdbcProfile]))
+      .flatMap { str =>
+        val formatted = Scalafmt.format(str, scalafmtConfig)
+        DBIO.from(Future.fromTry(formatted.toEither.toTry))
+      }
+
+  def writeToFileDBIO(
     baseDir: Path,
     slickConfig: Config,
     rules: GenerationRules
   )(implicit executionContext: ExecutionContext
+  ) =
+    codeStringFormatted(rules, slickConfig.getString("profile")).map { codeStr =>
+      val path = rules.filePath(baseDir)
+      Files.createDirectories(path.getParent)
+      Files.write(path, codeStr.getBytes())
+      path
+    }
+
+  def writeToFileSync(
+    baseDir: Path,
+    slickConfig: Config,
+    rules: GenerationRules,
+    timeout: Duration = Duration.Inf
+  )(implicit executionContext: ExecutionContext
   ): Path = {
-    val path = rules.filePath(baseDir)
-    val db   = JdbcBackend.Database.forConfig("", slickConfig)
-    val code =
-      try Await.result(db.run(codeString(rules, slickConfig.getString("profile"))), Duration.Inf)
-      finally db.close()
-    Files.createDirectories(path.getParent)
-    Files.write(path, code.getBytes())
+    val db = JdbcBackend.Database.forConfig("", slickConfig)
+
+    try
+      Await.result(db.run(writeToFileDBIO(baseDir, slickConfig, rules)), timeout)
+    finally
+      db.close()
   }
 }

slick-additions-codegen/src/test/resources/Models.scala

@@ -0,0 +1,31 @@
+package com.acme.models
+import slick.additions.entity.Lookup, io.circe.generic.JsonCodec,
+  monocle.macros.Lenses
+@JsonCodec @Lenses case class ColorsRow(name: String)
+object ColorsRow
+@JsonCodec @Lenses case class PeopleRow(
+    first: String,
+    last: String,
+    city: String = "New York",
+    dateJoined: java.time.LocalDate = java.time.LocalDate.now(),
+    balance: BigDecimal = BigDecimal("0.0"),
+    bestFriend: Option[Lookup[Long, PeopleRow]] = None,
+    col8: Option[Double] = None,
+    col9: Option[Boolean] = None,
+    col10: Option[Int] = None,
+    col11: Option[Int] = None,
+    col12: Option[Int] = None,
+    col13: Option[Int] = None,
+    col14: Option[Int] = None,
+    col15: Option[Int] = None,
+    col16: Option[Int] = None,
+    col17: Option[Int] = None,
+    col18: Option[Int] = None,
+    col19: Option[Int] = None,
+    col20: Option[Int] = None,
+    col21: Option[Int] = None,
+    col22: Option[Int] = None,
+    col23: Option[Int] = None,
+    col24: Option[Int] = None
+)
+object PeopleRow

slick-additions-codegen/src/test/resources/TableModules.scala

@@ -0,0 +1,166 @@
+package com.acme.tablemodules
+import slick.additions.AdditionsProfile
+import slick.lifted.MappedProjection
+trait SlickProfile extends slick.jdbc.H2Profile with AdditionsProfile {
+  object myApi extends JdbcAPI with AdditionsApi
+  override val api: myApi.type = myApi
+}
+object SlickProfile extends SlickProfile
+import SlickProfile.api._
+import slick.additions.entity.Lookup
+object TableModules {
+  object Colors extends EntityTableModule[Long, ColorsRow]("colors") {
+    class Row(tag: Tag) extends BaseEntRow(tag) {
+      override def keyColumnName = "id"
+      val name                   = column[String]("name")
+      def mapping: MappedProjection[ColorsRow] = name.<>(
+        {
+          ColorsRow.apply
+        },
+        ColorsRow.unapply
+      )
+    }
+  }
+  object People extends EntityTableModule[Long, PeopleRow]("people") {
+    class Row(tag: Tag) extends BaseEntRow(tag) {
+      override def keyColumnName = "id"
+      val first                  = column[String]("first")
+      val last                   = column[String]("last")
+      val city                   = column[String]("city")
+      val dateJoined             = column[java.time.LocalDate]("date_joined")
+      val balance                = column[BigDecimal]("balance")
+      val bestFriend = column[Option[Lookup[Long, PeopleRow]]]("best_friend")
+      val col8       = column[Option[Double]]("col8")
+      val col9       = column[Option[Boolean]]("col9")
+      val col10      = column[Option[Int]]("col10")
+      val col11      = column[Option[Int]]("col11")
+      val col12      = column[Option[Int]]("col12")
+      val col13      = column[Option[Int]]("col13")
+      val col14      = column[Option[Int]]("col14")
+      val col15      = column[Option[Int]]("col15")
+      val col16      = column[Option[Int]]("col16")
+      val col17      = column[Option[Int]]("col17")
+      val col18      = column[Option[Int]]("col18")
+      val col19      = column[Option[Int]]("col19")
+      val col20      = column[Option[Int]]("col20")
+      val col21      = column[Option[Int]]("col21")
+      val col22      = column[Option[Int]]("col22")
+      val col23      = column[Option[Int]]("col23")
+      val col24      = column[Option[Int]]("col24")
+      def mapping: MappedProjection[PeopleRow] = (
+        (
+          first,
+          last,
+          city,
+          dateJoined,
+          balance,
+          bestFriend,
+          col8,
+          col9,
+          col10,
+          col11,
+          col12,
+          col13,
+          col14,
+          col15,
+          col16,
+          col17,
+          col18,
+          col19,
+          col20,
+          col21,
+          col22,
+          col23
+        ),
+        col24
+      ).<>(
+        {
+          {
+            case (
+                  (
+                    first,
+                    last,
+                    city,
+                    dateJoined,
+                    balance,
+                    bestFriend,
+                    col8,
+                    col9,
+                    col10,
+                    col11,
+                    col12,
+                    col13,
+                    col14,
+                    col15,
+                    col16,
+                    col17,
+                    col18,
+                    col19,
+                    col20,
+                    col21,
+                    col22,
+                    col23
+                  ),
+                  col24
+                ) =>
+              PeopleRow(
+                first,
+                last,
+                city,
+                dateJoined,
+                balance,
+                bestFriend,
+                col8,
+                col9,
+                col10,
+                col11,
+                col12,
+                col13,
+                col14,
+                col15,
+                col16,
+                col17,
+                col18,
+                col19,
+                col20,
+                col21,
+                col22,
+                col23,
+                col24
+              )
+          }
+        },
+        (rec: PeopleRow) =>
+          Some(
+            (
+              (
+                rec.first,
+                rec.last,
+                rec.city,
+                rec.dateJoined,
+                rec.balance,
+                rec.bestFriend,
+                rec.col8,
+                rec.col9,
+                rec.col10,
+                rec.col11,
+                rec.col12,
+                rec.col13,
+                rec.col14,
+                rec.col15,
+                rec.col16,
+                rec.col17,
+                rec.col18,
+                rec.col19,
+                rec.col20,
+                rec.col21,
+                rec.col22,
+                rec.col23
+              ),
+              rec.col24
+            )
+          )
+      )
+    }
+  }
+}