Simplify

sake92 · Aug 16, 2024 · e322caa · e322caa
1 parent 8467fbd
commit e322caa
Show file tree

Hide file tree

Showing 3 changed files with 24 additions and 14 deletions.
diff --git a/generator/src/ba/sake/squery/generator/DbDefExtractor.scala b/generator/src/ba/sake/squery/generator/DbDefExtractor.scala
@@ -22,19 +22,16 @@ object DbDefExtractor {
 
 abstract class DbDefExtractor(ds: DataSource) {
 
-  def extract(schemaNames: Seq[String]): DbDef = Using.resource(ds.getConnection()) { connection =>
+  def extract(): DbDef = Using.resource(ds.getConnection()) { connection =>
     val databaseMetaData = connection.getMetaData()
     val dbName = databaseMetaData.getDatabaseProductName().toLowerCase
-    val existingSchemaNames = Using.resource(databaseMetaData.getSchemas()) { rs =>
+    val schemaNames = Using.resource(databaseMetaData.getSchemas()) { rs =>
       val buff = ArrayBuffer.empty[String]
       while (rs.next()) {
         buff += rs.getString("TABLE_SCHEM")
       }
       buff.toSeq
     }
-    val missingSchemaNames = schemaNames.toSet.diff(existingSchemaNames.toSet)
-    if (missingSchemaNames.nonEmpty)
-      throw new RuntimeException(s"Schemas do not exist in database: ${missingSchemaNames.mkString(", ")}")
     val schemaDefs = schemaNames.map { schemaName =>
       val tables = extractTables(connection, schemaName, databaseMetaData)
       SchemaDef(name = schemaName, tables = tables)

diff --git a/generator/src/ba/sake/squery/generator/SqueryGenerator.scala b/generator/src/ba/sake/squery/generator/SqueryGenerator.scala
@@ -1,17 +1,24 @@
 package ba.sake.squery.generator
 
+import java.io.File
+import javax.sql.DataSource
 import java.sql.{Array => _, _}
 import scala.util._
 import org.apache.commons.text.CaseUtils
 import com.typesafe.scalalogging.Logger
-import java.io.File
 
-class SqueryGenerator(config: SqueryGeneratorConfig = SqueryGeneratorConfig.Default) {
+class SqueryGenerator(ds: DataSource, config: SqueryGeneratorConfig = SqueryGeneratorConfig.Default) {
   private val logger = Logger(getClass.getName)
 
   private val Preamble = "/* DO NOT EDIT MANUALLY! Automatically generated by squery generator */"
 
-  def generateString(dbDef: DbDef, schemaNames: Seq[String]): String =
+  def generateString(schemaNames: Seq[String]): String = {
+    val extractor = DbDefExtractor(ds)
+    val dbDef = extractor.extract()
+    val missingSchemaNames = schemaNames.toSet.diff(dbDef.schemas.map(_.name).toSet)
+    if (missingSchemaNames.nonEmpty)
+      throw new RuntimeException(s"Schemas do not exist in database: ${missingSchemaNames.mkString(", ")}")
+
     schemaNames
       .map { schemaName =>
         dbDef.schemas.find(_.name == schemaName) match {
@@ -33,8 +40,15 @@ class SqueryGenerator(config: SqueryGeneratorConfig = SqueryGeneratorConfig.Defa
         }
       }
       .mkString("\n")
+  }
+
+  def generateFiles(schemaConfigs: Seq[SchemaConfig]): Unit = {
+    val extractor = DbDefExtractor(ds)
+    val dbDef = extractor.extract()
+    val missingSchemaNames = schemaConfigs.map(_.name).toSet.diff(dbDef.schemas.map(_.name).toSet)
+    if (missingSchemaNames.nonEmpty)
+      throw new RuntimeException(s"Schemas do not exist in database: ${missingSchemaNames.mkString(", ")}")
 
-  def generateFiles(dbDef: DbDef, schemaConfigs: Seq[SchemaConfig]): Unit =
     schemaConfigs.foreach { schemaConfig =>
       dbDef.schemas.find(_.name == schemaConfig.name) match {
         case Some(schemaDef) =>
@@ -77,6 +91,7 @@ class SqueryGenerator(config: SqueryGeneratorConfig = SqueryGeneratorConfig.Defa
           throw new RuntimeException(s"Schema '${schemaConfig.name}' does not exist")
       }
     }
+  }
 
   // (imports, models, repos)
   private def generateSchema(
@@ -150,7 +165,7 @@ class SqueryGenerator(config: SqueryGeneratorConfig = SqueryGeneratorConfig.Defa
       val caseClassName = transformName(tableDef.name, config.typeNameMapper) + config.rowTypeSuffix
       val daoClassName = transformName(tableDef.name, config.typeNameMapper) + "CrudDao"
       val findByIdWhereExpr = if (tableDef.pkColumns.isEmpty) {
-        // TODO throw or skip when no pks
+        // TODO skip when no pks
         ???
       } else if (tableDef.pkColumns.length == 1) {
         s"${tableDef.pkColumns.head.metadata.name} = $${id}"
@@ -160,6 +175,7 @@ class SqueryGenerator(config: SqueryGeneratorConfig = SqueryGeneratorConfig.Defa
             s"${pkCol.metadata.name} = $${id.${pkCol.metadata.name}}"
           }
           .mkString(" AND ")
+      // TODO create, update, delete
       val contents =
         s"""|trait ${daoClassName} {
             |  def countAll(): DbAction[Int] =

diff --git a/mill-plugin/src/ba/sake/squery/generator/mill/SqueryGeneratorModule.scala b/mill-plugin/src/ba/sake/squery/generator/mill/SqueryGeneratorModule.scala
@@ -54,11 +54,8 @@ trait SqueryGeneratorModule extends JavaModule {
         ds
       } else throw new RuntimeException(s"Unsupported database ${jdbcUrl}")
 
-    val extractor = DbDefExtractor(dataSource)
-    val dbDef = extractor.extract(squerySchemas().map(_._1))
-    val generator = new SqueryGenerator()
+    val generator = new SqueryGenerator(dataSource)
     generator.generateFiles(
-      dbDef,
       squerySchemas().map { case (schemaName, basePackage) =>
         SchemaConfig(
           name = schemaName,