Merge pull request #376 from victornguen/yugabytedb-wrapper

Add wrappers for YugabyteDB

build.sbt

@@ -100,6 +100,7 @@ lazy val root = (project in file("."))
     moduleRedpanda,
     moduleMinIO,
     moduleWireMock,
+    moduleYugabytedb,
     allOld
   )
   .settings(noPublishSettings)
@@ -525,3 +526,11 @@ lazy val moduleWireMock = (project in file("modules/wiremock"))
     name := "testcontainers-scala-wiremock",
     libraryDependencies ++= Dependencies.moduleWireMock.value
   )
+
+lazy val moduleYugabytedb = (project in file("modules/yugabytedb"))
+  .dependsOn(core % "compile->compile;test->test;provided->provided", scalatest % "test->test")
+  .settings(commonSettings)
+  .settings(
+    name := "testcontainers-scala-yugabytedb",
+    libraryDependencies ++= Dependencies.moduleYugabytedb.value
+  )

modules/yugabytedb/src/main/scala/com/dimafeng/testcontainers/YugabyteDBYCQLContainer.scala

@@ -0,0 +1,61 @@
+package com.dimafeng.testcontainers
+
+import org.testcontainers.containers.{YugabyteDBYCQLContainer => JavaYugabyteYCQLContainer}
+import org.testcontainers.utility.DockerImageName
+
+import java.net.InetSocketAddress
+
+class YugabyteDBYCQLContainer(
+                               underlying: JavaYugabyteYCQLContainer
+                             ) extends SingleContainer[JavaYugabyteYCQLContainer] {
+
+  override val container: JavaYugabyteYCQLContainer = underlying
+
+  def keyspace: String = container.getKeyspace
+
+  def localDc: String = container.getLocalDc
+
+  def username: String = container.getUsername
+
+  def password: String = container.getPassword
+
+  def contactPoint: InetSocketAddress = container.getContactPoint
+
+}
+
+object YugabyteDBYCQLContainer {
+
+  val defaultImage           = "yugabytedb/yugabyte"
+  val defaultTag             = "2.20.7.1-b10"
+  val defaultDockerImageName = s"$defaultImage:$defaultTag"
+
+  val ycqlPort    = 9042
+  val masterDashboardPort = 7000
+  val tserverDashboardPort = 9000
+
+  case class Def(
+                  dockerImageName: DockerImageName = DockerImageName.parse(YugabyteDBYCQLContainer.defaultDockerImageName),
+                  private val builder: List[JavaYugabyteYCQLContainer => JavaYugabyteYCQLContainer] = List.empty
+                ) extends ContainerDef {
+    override type Container = YugabyteDBYCQLContainer
+
+    def withKeyspaceName(keyspace: String): Def =
+      copy(builder = ((_: JavaYugabyteYCQLContainer).withKeyspaceName(keyspace)) :: builder)
+
+    def withUsername(username: String): Def =
+      copy(builder = ((_: JavaYugabyteYCQLContainer).withUsername(username)) :: builder)
+
+    def withPassword(password: String): Def =
+      copy(builder = ((_: JavaYugabyteYCQLContainer).withPassword(password)) :: builder)
+
+    def withInitScript(script: String): Def =
+      copy(builder = ((_: JavaYugabyteYCQLContainer).withInitScript(script)) :: builder)
+
+    override def createContainer(): YugabyteDBYCQLContainer = {
+      new YugabyteDBYCQLContainer(
+        builder
+          .foldRight(new JavaYugabyteYCQLContainer(dockerImageName))((f, underlying) => f(underlying))
+      )
+    }
+  }
+}

modules/yugabytedb/src/main/scala/com/dimafeng/testcontainers/YugabyteDBYSQLContainer.scala

@@ -0,0 +1,58 @@
+package com.dimafeng.testcontainers
+
+import org.testcontainers.containers.{YugabyteDBYSQLContainer => JavaYugabyteYSQLContainer}
+import org.testcontainers.utility.DockerImageName
+
+class YugabyteDBYSQLContainer(
+    underlying: JavaYugabyteYSQLContainer
+) extends SingleContainer[JavaYugabyteYSQLContainer] {
+
+  override val container: JavaYugabyteYSQLContainer = underlying
+
+  def databaseName: String = container.getDatabaseName
+
+  def username: String = container.getUsername
+
+  def password: String = container.getPassword
+
+  def driverClassName: String = container.getDriverClassName
+
+  def jdbcUrl: String = container.getJdbcUrl
+
+  def testQueryString: String = container.getTestQueryString
+
+}
+
+object YugabyteDBYSQLContainer {
+
+  val defaultImage           = "yugabytedb/yugabyte"
+  val defaultTag             = "2.20.7.1-b10"
+  val defaultDockerImageName = s"$defaultImage:$defaultTag"
+
+  val ysqlPort    = 5433
+  val masterDashboardPort = 7000
+  val tserverDashboardPort = 9000
+
+  case class Def(
+      dockerImageName: DockerImageName = DockerImageName.parse(YugabyteDBYSQLContainer.defaultDockerImageName),
+      private val builder: List[JavaYugabyteYSQLContainer => JavaYugabyteYSQLContainer] = List.empty
+  ) extends ContainerDef {
+    override type Container = YugabyteDBYSQLContainer
+
+    def withDatabaseName(database: String): Def =
+      copy(builder = ((_: JavaYugabyteYSQLContainer).withDatabaseName(database)) :: builder)
+
+    def withUsername(username: String): Def =
+      copy(builder = ((_: JavaYugabyteYSQLContainer).withUsername(username)) :: builder)
+
+    def withPassword(password: String): Def =
+      copy(builder = ((_: JavaYugabyteYSQLContainer).withPassword(password)) :: builder)
+
+    override def createContainer(): YugabyteDBYSQLContainer = {
+      new YugabyteDBYSQLContainer(
+        builder
+          .foldRight(new JavaYugabyteYSQLContainer(dockerImageName))((f, underlying) => f(underlying))
+      )
+    }
+  }
+}

modules/yugabytedb/src/test/resources/init_yql.sql

@@ -0,0 +1,7 @@
+CREATE TABLE dsql
+(
+    greet text primary key
+);
+
+INSERT INTO dsql (greet)
+VALUES ('Hello DSQL');

modules/yugabytedb/src/test/scala/com/dimafeng/testcontainers/YugabyteDBCQLSpec.scala

@@ -0,0 +1,52 @@
+package com.dimafeng.testcontainers
+
+import com.datastax.oss.driver.api.core.CqlSession
+import com.datastax.oss.driver.api.core.cql.ResultSet
+import com.dimafeng.testcontainers.scalatest.TestContainersForAll
+import org.scalatest.flatspec.AnyFlatSpec
+
+class YugabyteDBCQLSpec extends AnyFlatSpec with TestContainersForAll {
+  override type Containers = YugabyteDBYCQLContainer
+
+  val keyspace = "test_keyspace"
+
+  override def startContainers(): YugabyteDBYCQLContainer =
+    YugabyteDBYCQLContainer
+      .Def()
+      .withKeyspaceName(keyspace)
+      .withUsername("yugabyte")
+      .withPassword("yugabyte")
+      .withInitScript("init_yql.sql")
+      .start()
+
+  "Yugabytedb container" should "be started" in withContainers { yugabytedb =>
+    val result = YugabyteDBCQLSpec
+      .performQuery(yugabytedb, "SELECT release_version FROM system.local")
+
+    assert(result.wasApplied())
+  }
+
+  "Yugabytedb container" should "execute init script" in withContainers { yugabytedbContainer =>
+    val result = YugabyteDBCQLSpec
+      .performQuery(yugabytedbContainer, s"SELECT greet FROM $keyspace.dsql")
+
+    assert(
+      result.wasApplied() &&
+        result.one().getString(0) == "Hello DSQL"
+    )
+  }
+}
+
+object YugabyteDBCQLSpec {
+  private def performQuery(ycqlContainer: YugabyteDBYCQLContainer, cql: String): ResultSet = {
+    val session = CqlSession.builder
+      .withKeyspace(ycqlContainer.keyspace)
+      .withAuthCredentials(ycqlContainer.username, ycqlContainer.password)
+      .withLocalDatacenter(ycqlContainer.localDc)
+      .addContactPoint(ycqlContainer.contactPoint)
+      .build
+    try session.execute(cql)
+    finally if (session != null) session.close()
+  }
+
+}

modules/yugabytedb/src/test/scala/com/dimafeng/testcontainers/YugabyteDBSQLSpec.scala

@@ -0,0 +1,36 @@
+package com.dimafeng.testcontainers
+
+import com.dimafeng.testcontainers.scalatest.TestContainersForAll
+import org.scalatest.flatspec.AnyFlatSpec
+
+import java.sql.DriverManager
+
+class YugabyteDBSQLSpec extends AnyFlatSpec with TestContainersForAll {
+  override type Containers = YugabyteDBYSQLContainer
+
+  val databaseName = "test_db"
+
+  override def startContainers(): YugabyteDBYSQLContainer =
+    YugabyteDBYSQLContainer
+      .Def()
+      .withDatabaseName(databaseName)
+      .withUsername("yugabyte")
+      .withPassword("yugabyte")
+      .start()
+
+  "Yugabytedb container" should "be started" in withContainers { yugabytedb =>
+    Class.forName(yugabytedb.driverClassName)
+    val connection = DriverManager.getConnection(yugabytedb.jdbcUrl, yugabytedb.username, yugabytedb.password)
+
+    val preparedStatement = connection.prepareStatement(yugabytedb.testQueryString)
+    try {
+      val resultSet = preparedStatement.executeQuery()
+      resultSet.next()
+      assert(1 == resultSet.getInt(1))
+      resultSet.close()
+    } finally {
+      preparedStatement.close()
+      connection.close()
+    }
+  }
+}

project/Dependencies.scala

@@ -40,6 +40,8 @@ object Dependencies {
   private val jedisVersion = "5.0.0"
   private val wireMockTestcontainersVersion = "1.0-alpha-13"
   private val milvusSdkVersion = "2.4.1"
+  private val yugabyteJdbcVersion = "42.3.5-yb-6"
+  private val yugabyteJavaDriverVersion = "4.15.0-yb-2-TESTFIX.0"
 
   val allOld = Def.setting(
     PROVIDED(
@@ -342,4 +344,13 @@ object Dependencies {
       "com.softwaremill.sttp.client3" %% "core" % sttpVersion
     )
   )
+
+  val moduleYugabytedb = Def.setting(
+    COMPILE(
+      "org.testcontainers" % "yugabytedb" % testcontainersVersion
+    ) ++ TEST(
+      "com.yugabyte" % "jdbc-yugabytedb" % yugabyteJdbcVersion,
+      "com.yugabyte" % "java-driver-core" % yugabyteJavaDriverVersion,
+    )
+  )
 }