diff --git a/amber/src/test/scala/org/apache/texera/amber/engine/e2e/TestUtils.scala b/amber/src/test/scala/org/apache/texera/amber/engine/e2e/TestUtils.scala index ac71483a5d3..eae8bffdb3f 100644 --- a/amber/src/test/scala/org/apache/texera/amber/engine/e2e/TestUtils.scala +++ b/amber/src/test/scala/org/apache/texera/amber/engine/e2e/TestUtils.scala @@ -21,7 +21,6 @@ package org.apache.texera.amber.engine.e2e import com.twitter.util.{Await, Duration, Promise, Return, Throw, Try} import org.apache.pekko.actor.ActorSystem -import org.apache.texera.common.config.StorageConfig import org.apache.texera.amber.core.executor.OpExecInitInfo import org.apache.texera.amber.core.storage.DocumentFactory import org.apache.texera.amber.core.storage.model.VirtualDocument @@ -199,10 +198,28 @@ object TestUtils { * Note such test cases need to clean up the database at the end of running each test case. */ def initiateTexeraDBForTestCases(): Unit = { + org.apache.texera.dao.MockTexeraDB.ensureInitialized() + val embedded = org.apache.texera.dao.MockTexeraDB.getDBInstance + + val dbName = "texera_db_for_test_cases_" + java.util.UUID.randomUUID().toString.replace("-", "") + + scala.util.Using.resource(embedded.getPostgresDatabase.getConnection) { conn => + scala.util.Using.resource(conn.createStatement()) { stmt => + stmt.execute(s"DROP DATABASE IF EXISTS $dbName") + stmt.execute(s"CREATE DATABASE $dbName") + } + } + + val targetDbConn = embedded.getDatabase("postgres", dbName).getConnection + scala.util.Using.resource(targetDbConn.createStatement()) { stmt => + stmt.execute(org.apache.texera.dao.MockTexeraDB.getDDLScript) + } + targetDbConn.close() + SqlServer.initConnection( - StorageConfig.jdbcUrlForTestCases, - StorageConfig.jdbcUsername, - StorageConfig.jdbcPassword + embedded.getJdbcUrl("postgres", dbName), + "postgres", + "" ) } diff --git a/common/dao/src/test/scala/org/apache/texera/dao/MockTexeraDB.scala b/common/dao/src/test/scala/org/apache/texera/dao/MockTexeraDB.scala index 3ae97b62e03..bb6791c5f28 100644 --- a/common/dao/src/test/scala/org/apache/texera/dao/MockTexeraDB.scala +++ b/common/dao/src/test/scala/org/apache/texera/dao/MockTexeraDB.scala @@ -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 +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 - -trait MockTexeraDB { - - private var dbInstance: Option[EmbeddedPostgres] = None - private var dslContext: Option[DSLContext] = None - private val database: String = "texera_db" +import scala.util.Using + +/** + * Provides a JVM-singleton EmbeddedPostgres for tests. Multiple specs that mix + * 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 = "" - def executeScriptInJDBC(conn: Connection, script: String): Unit = { - assert(dbInstance.nonEmpty) - conn.prepareStatement(script).execute() - conn.close() - } + @volatile private var dbInstance: Option[EmbeddedPostgres] = None + @volatile private var ddlScript: Option[String] = None - def getDSLContext: DSLContext = { - dslContext match { - case Some(value) => value - case None => - throw new RuntimeException( - "test database is not initialized. Did you call initializeDBAndReplaceDSLContext()?" - ) - } - } + def ensureInitialized(): Unit = + synchronized { + if (dbInstance.isDefined && ddlScript.isDefined) return - def getDBInstance: EmbeddedPostgres = { - dbInstance match { - case Some(value) => value - case None => - throw new RuntimeException( - "test database is not initialized. Did you call initializeDBAndReplaceDSLContext()?" - ) - } - } + if (dbInstance.isEmpty) { + val driver = new org.postgresql.Driver() + DriverManager.registerDriver(driver) - def shutdownDB(): Unit = { - dbInstance match { - case Some(value) => - value.close() - dbInstance = None - dslContext = None - case None => - // do nothing - } - } + // 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() - def initializeDBAndReplaceDSLContext(): Unit = { - assert(dbInstance.isEmpty && dslContext.isEmpty) + val parts: Array[String] = content.split("(?m)^CREATE DATABASE :\"DB_NAME\";") + val sqlBody = if (parts.length > 1) parts(1) else content - val driver = new org.postgresql.Driver() - DriverManager.registerDriver(driver) + def removeCCommands(sql: String): String = + sql.linesIterator.filterNot(_.trim.startsWith("\\c")).mkString("\n") - val embedded = EmbeddedPostgres.builder().start() + var tablesAndIndexCreation = removeCCommands(sqlBody) - dbInstance = Some(embedded) + 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 - val ddlPath = { - Paths.get("sql/texera_ddl.sql").toRealPath() + // Cache the cleaned script so parallel suites don't have to re-read the file + ddlScript = Some(blockPattern.replaceAllIn(tablesAndIndexCreation, replacementText).trim) } - val source = Source.fromFile(ddlPath.toString) - val content = - try { - source.mkString - } finally { - source.close() + + 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) + 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", "") + SqlServer.getInstance().replaceDSLContext(scopedCtx) } - val parts: Array[String] = content.split("(?m)^CREATE DATABASE :\"DB_NAME\";") - def removeCCommands(sql: String): String = - sql.linesIterator - .filterNot(_.trim.startsWith("\\c")) - .mkString("\n") - val createDBStatement = - """DROP DATABASE IF EXISTS texera_db; - |CREATE DATABASE texera_db;""".stripMargin - executeScriptInJDBC(embedded.getPostgresDatabase.getConnection, createDBStatement) - val texeraDB = embedded.getDatabase(username, database) - var tablesAndIndexCreation = removeCCommands(parts(1)) - - // remove indexes creation for pgroonga because we cannot install the plugin - 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 - // replace with native fulltext indexes - 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 - - tablesAndIndexCreation = blockPattern.replaceAllIn(tablesAndIndexCreation, replacementText).trim - executeScriptInJDBC(texeraDB.getConnection, tablesAndIndexCreation) - SqlServer.initConnection(embedded.getJdbcUrl(username, database), username, password) + } + + abstract override def withFixture(test: NoArgTest): Outcome = { + initializeDBAndReplaceDSLContext() + val sqlServerInstance = SqlServer.getInstance() - dslContext = Some(DSL.using(texeraDB, SQLDialect.POSTGRES)) + val activeContext = testScopedContext.get - sqlServerInstance.replaceDSLContext(dslContext.get) + try { + sqlServerInstance.replaceDSLContext(activeContext) + 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 + EXECUTE 'TRUNCATE TABLE ' || quote_ident(r.tablename) || ' CASCADE'; + END LOOP; + END $$; + """ + ) + } + } + } catch { + case e: Exception => e.printStackTrace() + } + } } + + def getDSLContext: DSLContext = + synchronized { + if (testScopedContext.isEmpty) { + initializeDBAndReplaceDSLContext() + } + testScopedContext.get + } + + def getDBInstance: EmbeddedPostgres = MockTexeraDB.getDBInstance + + def shutdownDB(): Unit = + synchronized { + try dataSource.foreach(ds => if (!ds.isClosed) ds.close()) + catch { case e: Exception => e.printStackTrace() } + finally { dataSource = None; testScopedContext = None } + } } diff --git a/file-service/src/test/scala/org/apache/texera/service/resource/DatasetResourceSpec.scala b/file-service/src/test/scala/org/apache/texera/service/resource/DatasetResourceSpec.scala index 1730d12a0a0..50088abc161 100644 --- a/file-service/src/test/scala/org/apache/texera/service/resource/DatasetResourceSpec.scala +++ b/file-service/src/test/scala/org/apache/texera/service/resource/DatasetResourceSpec.scala @@ -1233,8 +1233,7 @@ class DatasetResourceSpec val filePath = uniqueFilePath("init-session-row-locked") initUpload(filePath, numParts = 2).getStatus shouldEqual 200 - val connectionProvider = getDSLContext.configuration().connectionProvider() - val connection = connectionProvider.acquire() + val connection = MockTexeraDB.getDBInstance.getDatabase("postgres", uniqueDbName).getConnection connection.setAutoCommit(false) try { @@ -1256,7 +1255,7 @@ class DatasetResourceSpec ex.getResponse.getStatus shouldEqual 409 } finally { connection.rollback() - connectionProvider.release(connection) + connection.close() } // lock released => init works again @@ -1622,8 +1621,8 @@ class DatasetResourceSpec val filePath = uniqueFilePath("init-lock-409") initUpload(filePath, numParts = 2).getStatus shouldEqual 200 - val connectionProvider = getDSLContext.configuration().connectionProvider() - val connection = connectionProvider.acquire() + // Open a completely independent connection to simulate a second concurrent user + val connection = MockTexeraDB.getDBInstance.getDatabase("postgres", uniqueDbName).getConnection connection.setAutoCommit(false) try { @@ -1645,7 +1644,7 @@ class DatasetResourceSpec ex.getResponse.getStatus shouldEqual 409 } finally { connection.rollback() - connectionProvider.release(connection) + connection.close() } } @@ -1865,8 +1864,7 @@ class DatasetResourceSpec initUpload(filePath, numParts = 2) val uploadId = fetchUploadIdOrFail(filePath) - val connectionProvider = getDSLContext.configuration().connectionProvider() - val connection = connectionProvider.acquire() + val connection = MockTexeraDB.getDBInstance.getDatabase("postgres", uniqueDbName).getConnection connection.setAutoCommit(false) try { @@ -1887,7 +1885,7 @@ class DatasetResourceSpec assertStatus(ex, 409) } finally { connection.rollback() - connectionProvider.release(connection) + connection.close() } uploadPart(filePath, 1, minPartBytes(3.toByte)).getStatus shouldEqual 200 @@ -1898,8 +1896,7 @@ class DatasetResourceSpec initUpload(filePath, numParts = 2) val uploadId = fetchUploadIdOrFail(filePath) - val connectionProvider = getDSLContext.configuration().connectionProvider() - val connection = connectionProvider.acquire() + val connection = MockTexeraDB.getDBInstance.getDatabase("postgres", uniqueDbName).getConnection connection.setAutoCommit(false) try { @@ -1917,7 +1914,7 @@ class DatasetResourceSpec uploadPart(filePath, 2, tinyBytes(9.toByte)).getStatus shouldEqual 200 } finally { connection.rollback() - connectionProvider.release(connection) + connection.close() } } @@ -2104,8 +2101,7 @@ class DatasetResourceSpec initUpload(filePath, numParts = 1) uploadPart(filePath, 1, tinyBytes(1.toByte)).getStatus shouldEqual 200 - val connectionProvider = getDSLContext.configuration().connectionProvider() - val connection = connectionProvider.acquire() + val connection = MockTexeraDB.getDBInstance.getDatabase("postgres", uniqueDbName).getConnection connection.setAutoCommit(false) try { @@ -2125,7 +2121,7 @@ class DatasetResourceSpec assertStatus(ex, 409) } finally { connection.rollback() - connectionProvider.release(connection) + connection.close() } } @@ -2165,8 +2161,7 @@ class DatasetResourceSpec val filePath = uniqueFilePath("abort-lock-race") initUpload(filePath, numParts = 1) - val connectionProvider = getDSLContext.configuration().connectionProvider() - val connection = connectionProvider.acquire() + val connection = MockTexeraDB.getDBInstance.getDatabase("postgres", uniqueDbName).getConnection connection.setAutoCommit(false) try { @@ -2186,7 +2181,7 @@ class DatasetResourceSpec assertStatus(ex, 409) } finally { connection.rollback() - connectionProvider.release(connection) + connection.close() } }