refactor: Refactor MockTexeraDB into a JVM singleton #6238
There are no files selected for viewing
| Original file line number | Diff line number | Diff line change |
|---|---|---|
|
|
@@ -19,146 +19,173 @@ | |
|
|
||
| package org.apache.texera.dao | ||
|
|
||
| import com.zaxxer.hikari.{HikariConfig, HikariDataSource} | ||
| import io.zonky.test.db.postgres.embedded.EmbeddedPostgres | ||
| 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.DriverManager | ||
| import scala.io.Source | ||
| import scala.util.Using | ||
|
|
||
| /** | ||
| * Provides a JVM-singleton EmbeddedPostgres for tests. Multiple specs that mix | ||
|
Contributor
There was a problem hiding this comment. This credits the singleton with avoiding |
||
| * in this trait share one Postgres instance for the lifetime of the JVM, which | ||
| * avoids the OverlappingFileLockException that occurs when each spec tries to | ||
| * extract the embedded Postgres binaries into the same directory in parallel. | ||
| */ | ||
| object MockTexeraDB { | ||
| private val username: String = "postgres" | ||
| private val password: String = "" | ||
|
|
||
| @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) | ||
|
|
||
| // Boot the heavy JVM engine exactly once | ||
| dbInstance = Some(EmbeddedPostgres.builder().start()) | ||
| } | ||
|
|
||
| val ddlPath = Paths.get("sql/texera_ddl.sql").toRealPath() | ||
| val source = Source.fromFile(ddlPath.toString) | ||
| val content = | ||
| try source.mkString | ||
| finally source.close() | ||
|
|
||
| val parts: Array[String] = content.split("(?m)^CREATE DATABASE :\"DB_NAME\";") | ||
| val sqlBody = if (parts.length > 1) parts(1) else content | ||
|
Contributor
There was a problem hiding this comment. If the |
||
|
|
||
| def removeCCommands(sql: String): String = | ||
| sql.linesIterator.filterNot(_.trim.startsWith("\\c")).mkString("\n") | ||
|
|
||
| var tablesAndIndexCreation = removeCCommands(sqlBody) | ||
|
Contributor
There was a problem hiding this comment.
|
||
|
|
||
| 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 getDBInstance: EmbeddedPostgres = | ||
| dbInstance.getOrElse(throw new RuntimeException("DB not initialized")) | ||
| def getDDLScript: String = ddlScript.getOrElse(throw new RuntimeException("DDL not loaded")) | ||
| } | ||
|
|
||
| trait MockTexeraDB extends TestSuiteMixin { this: TestSuite => | ||
| private var testScopedContext: Option[DSLContext] = None | ||
| protected var dataSource: Option[HikariDataSource] = None | ||
| protected var uniqueDbName: String = "" | ||
|
|
||
| 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) | ||
|
|
||
| // Back the test DSLContext with a real HikariCP pool so that concurrent | ||
| // transactions acquire *distinct* connections (matching production), rather | ||
| // than trampling one shared connection's autoCommit flag. | ||
| val hikariConfig = new HikariConfig() | ||
| hikariConfig.setJdbcUrl(jdbcUrl) | ||
| hikariConfig.setUsername("postgres") | ||
| hikariConfig.setPassword("") | ||
| // Must exceed the maximum number of concurrent test threads. | ||
| hikariConfig.setMaximumPoolSize(10) | ||
|
Contributor
There was a problem hiding this comment.
|
||
| val ds = new HikariDataSource(hikariConfig) | ||
| dataSource = Some(ds) | ||
|
|
||
| val jooqCfg = new DefaultConfiguration() | ||
| jooqCfg.set(new DataSourceConnectionProvider(ds)) | ||
| jooqCfg.set(SQLDialect.POSTGRES) | ||
| val scopedCtx = DSL.using(jooqCfg) | ||
| testScopedContext = Some(scopedCtx) | ||
|
|
||
| // Point the Texera backend exactly to this suite's isolated database | ||
| SqlServer.initConnection(jdbcUrl, "postgres", "") | ||
|
Contributor
There was a problem hiding this comment.
|
||
| SqlServer.getInstance().replaceDSLContext(scopedCtx) | ||
| } | ||
| } | ||
|
|
||
| abstract override def withFixture(test: NoArgTest): Outcome = { | ||
| initializeDBAndReplaceDSLContext() | ||
|
|
||
| val sqlServerInstance = SqlServer.getInstance() | ||
| val activeContext = testScopedContext.get | ||
|
|
||
| try { | ||
| sqlServerInstance.replaceDSLContext(activeContext) | ||
|
Contributor
There was a problem hiding this comment. Per-suite isolation is real for |
||
| super.withFixture(test) | ||
| } finally { | ||
| try { | ||
| // Truncate all tables on a pooled connection so each test starts clean. | ||
| Using.resource(dataSource.get.getConnection) { conn => | ||
| Using.resource(conn.createStatement()) { stmt => | ||
| stmt.execute( | ||
| """ | ||
| DO $$ DECLARE | ||
| r RECORD; | ||
| BEGIN | ||
| FOR r IN (SELECT tablename FROM pg_tables WHERE schemaname = 'public') LOOP | ||
|
Contributor
There was a problem hiding this comment. This per-test cleanup is a silent no-op (targets the wrong schema). It truncates the
Contributor
Author
There was a problem hiding this comment. I think removing the truncation for now and leaving the truncation logic for later would be best? This allows us to get the speed-up now without ballooning the responsibilities of this PR? |
||
| EXECUTE 'TRUNCATE TABLE ' || quote_ident(r.tablename) || ' CASCADE'; | ||
| END LOOP; | ||
| END $$; | ||
| """ | ||
| ) | ||
| } | ||
| } | ||
| } catch { | ||
| case e: Exception => e.printStackTrace() | ||
|
Contributor
There was a problem hiding this comment. If the truncation throws it's only printed and swallowed, so the next test can start against dirty state. If the cleanup stays, a failure to clean should fail loudly rather than print-and-continue. |
||
| } | ||
| } | ||
| } | ||
|
|
||
| def getDSLContext: DSLContext = | ||
| synchronized { | ||
| if (testScopedContext.isEmpty) { | ||
| initializeDBAndReplaceDSLContext() | ||
| } | ||
| testScopedContext.get | ||
| } | ||
|
|
||
| def getDBInstance: EmbeddedPostgres = MockTexeraDB.getDBInstance | ||
|
Contributor
There was a problem hiding this comment. Minor readability: |
||
|
|
||
| def shutdownDB(): Unit = | ||
|
Contributor
There was a problem hiding this comment. Two small things: the name no longer matches its behavior (it closes the per-suite pool, not a DB), and the created databases are never dropped (fine now, matters once suites run in parallel). Also #6063 asked to no-op |
||
| synchronized { | ||
| try dataSource.foreach(ds => if (!ds.isClosed) ds.close()) | ||
| catch { case e: Exception => e.printStackTrace() } | ||
| finally { dataSource = None; testScopedContext = None } | ||
| } | ||
| } | ||
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
A few things in this method: (1) it duplicates the "create DB + run DDL" sequence from the trait (
MockTexeraDB.scala:101-114) and the copies have already drifted — worth one shared helper onobject MockTexeraDB. (2)targetDbConnis closed manually after theUsing.resource, so it leaks ifexecutethrows — wrap it inUsing.resourcetoo. (3) TheDROP DATABASE IF EXISTSon a freshly generated UUID name (line 208) is dead. (4) This quietly moves the e2e specs off the external test Postgres onto the embedded one — a real (good) change worth a line in the description; it depends on #4179. I ran the e2e specs this way on JDK 11 (46 tests ×3, noConnection refused), so the switch looks safe for sequential runs.