Ma77Ball commented on code in PR #6238:
URL: https://github.com/apache/texera/pull/6238#discussion_r3610276011


##########
common/dao/src/test/scala/org/apache/texera/dao/MockTexeraDB.scala:
##########
@@ -19,146 +19,175 @@
 
 package org.apache.texera.dao
 
+import com.zaxxer.hikari.{HikariConfig, HikariDataSource}
 import io.zonky.test.db.postgres.embedded.EmbeddedPostgres
-import org.jooq.impl.DSL
+import org.apache.texera.dao.MockTexeraDB.{MaxPoolSize, password, username}
+import org.jooq.impl.{DSL, DataSourceConnectionProvider, DefaultConfiguration}
 import org.jooq.{DSLContext, SQLDialect}
+import org.scalatest.{Outcome, TestSuite, TestSuiteMixin}
 
 import java.nio.file.Paths
-import java.sql.{Connection, DriverManager}
+import java.sql.DriverManager
 import scala.io.Source
+import scala.util.Using
 
-trait MockTexeraDB {
-
-  private var dbInstance: Option[EmbeddedPostgres] = None
-  private var dslContext: Option[DSLContext] = None
-  private val database: String = "texera_db"
+/**
+  * Provides a JVM-singleton EmbeddedPostgres for tests. Multiple specs that 
mix
+  * in this trait share one Postgres instance for the lifetime of the JVM.
+  */
+object MockTexeraDB {
   private val username: String = "postgres"
   private val password: String = ""
+  private val texeraDDLPath = "sql/texera_ddl.sql"
+  private val splitDatabaseRegex = "(?m)^CREATE DATABASE :\"DB_NAME\";"
 
-  def executeScriptInJDBC(conn: Connection, script: String): Unit = {
-    assert(dbInstance.nonEmpty)
-    conn.prepareStatement(script).execute()
-    conn.close()
-  }
+  val MaxPoolSize: Int = math.max(10, Runtime.getRuntime.availableProcessors() 
* 2)
 
-  def getDSLContext: DSLContext = {
-    dslContext match {
-      case Some(value) => value
-      case None =>
-        throw new RuntimeException(
-          "test database is not initialized. Did you call 
initializeDBAndReplaceDSLContext()?"
-        )
-    }
-  }
+  @volatile private var dbInstance: Option[EmbeddedPostgres] = None
+  @volatile private var ddlScript: Option[String] = None
+
+  def ensureInitialized(): Unit =
+    synchronized {
+      if (dbInstance.isDefined && ddlScript.isDefined) return
+
+      if (dbInstance.isEmpty) {
+        val driver = new org.postgresql.Driver()
+        DriverManager.registerDriver(driver)
 
-  def getDBInstance: EmbeddedPostgres = {
-    dbInstance match {
-      case Some(value) => value
-      case None =>
-        throw new RuntimeException(
-          "test database is not initialized. Did you call 
initializeDBAndReplaceDSLContext()?"
+        // Boot the heavy JVM engine exactly once
+        dbInstance = Some(EmbeddedPostgres.builder().start())
+      }
+
+      val ddlPath = Paths.get(texeraDDLPath).toRealPath()
+      val source = Source.fromFile(ddlPath.toString)
+      val content =
+        try source.mkString
+        finally source.close()
+
+      val parts: Array[String] = content.split(splitDatabaseRegex)
+      val sqlBody = parts
+        .lift(1)
+        .getOrElse(
+          throw new RuntimeException(
+            s"Couldn't split SQL body from $texeraDDLPath: " +
+              s"expected it to match pattern $splitDatabaseRegex"
+          )
         )
-    }
-  }
 
-  def shutdownDB(): Unit = {
-    dbInstance match {
-      case Some(value) =>
-        value.close()
-        dbInstance = None
-        dslContext = None
-      case None =>
-      // do nothing
+      def removeCCommands(sql: String): String =
+        sql.linesIterator.filterNot(_.trim.startsWith("\\c")).mkString("\n")
+
+      val tablesAndIndexCreation = removeCCommands(sqlBody)
+
+      val blockPattern =
+        """(?s)-- START Fulltext search index creation \(DO NOT EDIT THIS 
LINE\).*?-- END Fulltext search index creation \(DO NOT EDIT THIS LINE\)\n?""".r
+      val replacementText =
+        """CREATE INDEX idx_workflow_name_description_content ON workflow 
USING GIN (to_tsvector('english', COALESCE(name, '') || ' ' || 
COALESCE(description, '') || ' ' || COALESCE(content, '')));
+        |CREATE INDEX idx_user_name ON "user" USING GIN 
(to_tsvector('english', COALESCE(name, '')));
+        |CREATE INDEX idx_user_project_name_description ON project USING GIN 
(to_tsvector('english', COALESCE(name, '') || ' ' || COALESCE(description, 
'')));
+        |CREATE INDEX idx_dataset_name_description ON dataset USING GIN 
(to_tsvector('english', COALESCE(name, '') || ' ' || COALESCE(description, 
'')));
+        |CREATE INDEX idx_dataset_version_name ON dataset_version USING GIN 
(to_tsvector('english', COALESCE(name, '')));""".stripMargin
+
+      // Cache the cleaned script so parallel suites don't have to re-read the 
file
+      ddlScript = Some(blockPattern.replaceAllIn(tablesAndIndexCreation, 
replacementText).trim)
     }
-  }
 
-  def initializeDBAndReplaceDSLContext(): Unit = {
-    assert(dbInstance.isEmpty && dslContext.isEmpty)
+  def getDBInstance: EmbeddedPostgres =
+    dbInstance.getOrElse(throw new RuntimeException("DB not initialized"))
+  def getDDLScript: String = ddlScript.getOrElse(throw new 
RuntimeException("DDL not loaded"))
+}
 
-    val driver = new org.postgresql.Driver()
-    DriverManager.registerDriver(driver)
+trait MockTexeraDB extends TestSuiteMixin { this: TestSuite =>
+  private var testScopedContext: Option[DSLContext] = None
+  protected var dataSource: Option[HikariDataSource] = None
+  protected var uniqueDbName: String = ""
+
+  def createHikariConfig(jbdcUrl: String): HikariConfig = {
+    val hikariConfig = new HikariConfig()
+    hikariConfig.setJdbcUrl(jbdcUrl)
+    hikariConfig.setUsername(username)
+    hikariConfig.setPassword(password)
+    hikariConfig.setMaximumPoolSize(MaxPoolSize)
+    hikariConfig
+  }
 
-    val embedded = EmbeddedPostgres.builder().start()
+  def initializeDBAndReplaceDSLContext(): Unit =
+    synchronized {
+      if (dataSource.isEmpty || dataSource.get.isClosed) {
+        MockTexeraDB.ensureInitialized()
+        val embedded = MockTexeraDB.getDBInstance
+
+        uniqueDbName = "texera_db_" + 
java.util.UUID.randomUUID().toString.replace("-", "")
+        Using.resource(embedded.getPostgresDatabase.getConnection) { 
defaultConn =>
+          Using.resource(defaultConn.createStatement()) { stmt =>
+            stmt.execute(s"CREATE DATABASE $uniqueDbName")
+          }
+        }
+
+        // Run the DDL once via a throwaway connection (autoCommit is TRUE by 
default,
+        // so the schema is permanently committed to this suite's isolated 
database).
+        Using.resource(embedded.getDatabase("postgres", 
uniqueDbName).getConnection) { conn =>
+          Using.resource(conn.createStatement()) { stmt =>
+            stmt.execute(MockTexeraDB.getDDLScript)
+          }
+        }
+
+        val jdbcUrl = embedded.getJdbcUrl("postgres", uniqueDbName)
+        val ds = new HikariDataSource(createHikariConfig(jbdcUrl = jdbcUrl))
+        dataSource = Some(ds)
+
+        val jooqCfg = new DefaultConfiguration()
+        jooqCfg.set(new DataSourceConnectionProvider(ds))
+        jooqCfg.set(SQLDialect.POSTGRES)
+        val scopedCtx = DSL.using(jooqCfg)
+        testScopedContext = Some(scopedCtx)
+
+        SqlServer.initConnection(jdbcUrl, username, password)
+        SqlServer.getInstance().replaceDSLContext(scopedCtx)
+      }
+    }
 
-    dbInstance = Some(embedded)
+  abstract override def withFixture(test: NoArgTest): Outcome = {

Review Comment:
   There should be a comment here that highlights that this design is only safe 
when test suites are run sequentially. Since the specs share one JVM Postgres, 
enabling parallel execution in the future will cause issues (it would be good 
to avoid an expensive debug session). 



-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: [email protected]

For queries about this service, please contact Infrastructure at:
[email protected]

Reply via email to