Add TiDB module (#5511)

Cheese · bsideup · kiview · web-flow · commit 54d7f3931d88 · 2022-08-15T16:01:46.000+02:00
Co-authored-by: Sergei Egorov &lt;bsideup@gmail.com&gt;
Co-authored-by: Kevin Wittek &lt;kevin@wittek.dev&gt;
Co-authored-by: Eddú Meléndez Gonzales &lt;eddu.melendez@gmail.com&gt;
diff --git a/.github/ISSUE_TEMPLATE/bug_report.yaml b/.github/ISSUE_TEMPLATE/bug_report.yaml
@@ -43,6 +43,7 @@ body:
         - RabbitMQ
         - Selenium
         - Solr
+        - TiDB
         - ToxiProxy
         - Trino
         - Vault
diff --git a/.github/ISSUE_TEMPLATE/enhancement.yaml b/.github/ISSUE_TEMPLATE/enhancement.yaml
@@ -43,6 +43,7 @@ body:
         - RabbitMQ
         - Selenium
         - Solr
+        - TiDB
         - ToxiProxy
         - Trino
         - Vault
diff --git a/.github/ISSUE_TEMPLATE/feature.yaml b/.github/ISSUE_TEMPLATE/feature.yaml
@@ -43,6 +43,7 @@ body:
         - RabbitMQ
         - Selenium
         - Solr
+        - TiDB
         - ToxiProxy
         - Trino
         - Vault
diff --git a/.github/dependabot.yml b/.github/dependabot.yml
@@ -197,6 +197,11 @@ updates:
     schedule:
       interval: "monthly"
     open-pull-requests-limit: 10
+  - package-ecosystem: "gradle"
+    directory: "/modules/tidb"
+    schedule:
+      interval: "monthly"
+    open-pull-requests-limit: 10
   - package-ecosystem: "gradle"
     directory: "/modules/toxiproxy"
     schedule:
diff --git a/.github/labeler.yml b/.github/labeler.yml
@@ -73,6 +73,8 @@
   - modules/solr/*
 "modules/spock":
   - modules/spock/*
+"modules/tidb":
+  - modules/tidb/*
 "modules/toxiproxy":
   - modules/toxiproxy/*
 "modules/trino":
diff --git a/docs/modules/databases/jdbc.md b/docs/modules/databases/jdbc.md
@@ -51,6 +51,10 @@ Insert `tc:` after `jdbc:` as follows. Note that the hostname, port and database
 
 `jdbc:tc:cockroach:v21.2.3:///databasename`
 
+#### Using TiDB
+
+`jdbc:tc:tidb:v6.1.0:///databasename`
+
 ### Using a classpath init script
 
 Testcontainers can run an init script after the database container is started, but before your code is given a connection to it. The script must be on the classpath, and is referenced as follows:
diff --git a/docs/modules/databases/tidb.md b/docs/modules/databases/tidb.md
@@ -0,0 +1,25 @@
+# TiDB Module
+
+See [Database containers](./index.md) for documentation and usage that is common to all relational database container types.
+
+## Adding this module to your project dependencies
+
+Add the following dependency to your `pom.xml`/`build.gradle` file:
+
+=== "Gradle"
+    ```groovy
+    testImplementation "org.testcontainers:tidb:{{latest_version}}"
+    ```
+
+=== "Maven"
+    ```xml
+    <dependency>
+        <groupId>org.testcontainers</groupId>
+        <artifactId>tidb</artifactId>
+        <version>{{latest_version}}</version>
+        <scope>test</scope>
+    </dependency>
+    ```
+
+!!! hint
+    Adding this Testcontainers library JAR will not automatically add a database driver JAR to your project. You should ensure that your project also has a suitable database driver as a dependency.
diff --git a/mkdocs.yml b/mkdocs.yml
@@ -63,6 +63,7 @@ nav:
                 - modules/databases/orientdb.md
                 - modules/databases/postgres.md
                 - modules/databases/presto.md
+                - modules/databases/tidb.md
                 - modules/databases/trino.md
           - modules/azure.md
           - modules/docker_compose.md
diff --git a/modules/tidb/build.gradle b/modules/tidb/build.gradle
@@ -0,0 +1,10 @@
+description = "Testcontainers :: JDBC :: TiDB"
+
+dependencies {
+    api project(':jdbc')
+
+    testImplementation project(':jdbc-test')
+    testImplementation 'mysql:mysql-connector-java:8.0.29'
+
+    compileOnly 'org.jetbrains:annotations:23.0.0'
+}
diff --git a/modules/tidb/sql/init_mysql.sql b/modules/tidb/sql/init_mysql.sql
@@ -0,0 +1,5 @@
+CREATE TABLE bar (
+  foo VARCHAR(255)
+);
+
+INSERT INTO bar (foo) VALUES ('hello world');
diff --git a/modules/tidb/src/main/java/org/testcontainers/tidb/TiDBContainer.java b/modules/tidb/src/main/java/org/testcontainers/tidb/TiDBContainer.java
@@ -0,0 +1,126 @@
+package org.testcontainers.tidb;
+
+import org.jetbrains.annotations.NotNull;
+import org.testcontainers.containers.JdbcDatabaseContainer;
+import org.testcontainers.containers.wait.strategy.HttpWaitStrategy;
+import org.testcontainers.utility.DockerImageName;
+
+import java.time.Duration;
+import java.util.HashSet;
+import java.util.Set;
+
+/**
+ * Testcontainers implementation for TiDB.
+ *
+ * @author Icemap
+ */
+public class TiDBContainer extends JdbcDatabaseContainer<TiDBContainer> {
+
+    static final String NAME = "tidb";
+
+    static final String DOCKER_IMAGE_NAME = "pingcap/tidb";
+
+    private static final DockerImageName DEFAULT_IMAGE_NAME = DockerImageName.parse(DOCKER_IMAGE_NAME);
+
+    private static final Integer TIDB_PORT = 4000;
+
+    private static final int REST_API_PORT = 10080;
+
+    private String databaseName = "test";
+
+    private String username = "root";
+
+    private String password = "";
+
+    public TiDBContainer(String dockerImageName) {
+        this(DockerImageName.parse(dockerImageName));
+    }
+
+    public TiDBContainer(final DockerImageName dockerImageName) {
+        super(dockerImageName);
+        dockerImageName.assertCompatibleWith(DEFAULT_IMAGE_NAME);
+
+        addExposedPorts(TIDB_PORT, REST_API_PORT);
+
+        waitingFor(
+            new HttpWaitStrategy()
+                .forPath("/status")
+                .forPort(REST_API_PORT)
+                .forStatusCode(200)
+                .withStartupTimeout(Duration.ofMinutes(1))
+        );
+    }
+
+    @NotNull
+    @Override
+    protected Set<Integer> getLivenessCheckPorts() {
+        return new HashSet<>(getMappedPort(TIDB_PORT));
+    }
+
+    @Override
+    public String getDriverClassName() {
+        try {
+            Class.forName("com.mysql.cj.jdbc.Driver");
+            return "com.mysql.cj.jdbc.Driver";
+        } catch (ClassNotFoundException e) {
+            return "com.mysql.jdbc.Driver";
+        }
+    }
+
+    @Override
+    public String getJdbcUrl() {
+        String additionalUrlParams = constructUrlParameters("?", "&");
+        return "jdbc:mysql://" + getHost() + ":" + getMappedPort(TIDB_PORT) + "/" + databaseName + additionalUrlParams;
+    }
+
+    @Override
+    protected String constructUrlForConnection(String queryString) {
+        String url = super.constructUrlForConnection(queryString);
+
+        if (!url.contains("useSSL=")) {
+            String separator = url.contains("?") ? "&" : "?";
+            url = url + separator + "useSSL=false";
+        }
+
+        if (!url.contains("allowPublicKeyRetrieval=")) {
+            url = url + "&allowPublicKeyRetrieval=true";
+        }
+
+        return url;
+    }
+
+    @Override
+    public String getDatabaseName() {
+        return databaseName;
+    }
+
+    @Override
+    public String getUsername() {
+        return username;
+    }
+
+    @Override
+    public String getPassword() {
+        return password;
+    }
+
+    @Override
+    public String getTestQueryString() {
+        return "SELECT 1";
+    }
+
+    @Override
+    public TiDBContainer withDatabaseName(final String databaseName) {
+        throw new UnsupportedOperationException("The TiDB docker image does not currently support this");
+    }
+
+    @Override
+    public TiDBContainer withUsername(final String username) {
+        throw new UnsupportedOperationException("The TiDB docker image does not currently support this");
+    }
+
+    @Override
+    public TiDBContainer withPassword(final String password) {
+        throw new UnsupportedOperationException("The TiDB docker image does not currently support this");
+    }
+}
diff --git a/modules/tidb/src/main/java/org/testcontainers/tidb/TiDBContainerProvider.java b/modules/tidb/src/main/java/org/testcontainers/tidb/TiDBContainerProvider.java
@@ -0,0 +1,32 @@
+package org.testcontainers.tidb;
+
+import org.testcontainers.containers.JdbcDatabaseContainer;
+import org.testcontainers.containers.JdbcDatabaseContainerProvider;
+import org.testcontainers.utility.DockerImageName;
+
+/**
+ * Factory for TiDB containers.
+ */
+public class TiDBContainerProvider extends JdbcDatabaseContainerProvider {
+
+    private static final String DEFAULT_TAG = "v6.1.0";
+
+    @Override
+    public boolean supports(String databaseType) {
+        return databaseType.equals(TiDBContainer.NAME);
+    }
+
+    @Override
+    public JdbcDatabaseContainer newInstance() {
+        return newInstance(DEFAULT_TAG);
+    }
+
+    @Override
+    public JdbcDatabaseContainer newInstance(String tag) {
+        if (tag != null) {
+            return new TiDBContainer(DockerImageName.parse(TiDBContainer.DOCKER_IMAGE_NAME).withTag(tag));
+        } else {
+            return newInstance();
+        }
+    }
+}
diff --git a/modules/tidb/src/main/resources/META-INF/services/org.testcontainers.containers.JdbcDatabaseContainerProvider b/modules/tidb/src/main/resources/META-INF/services/org.testcontainers.containers.JdbcDatabaseContainerProvider
@@ -0,0 +1 @@
+org.testcontainers.tidb.TiDBContainerProvider
diff --git a/modules/tidb/src/test/java/org/testcontainers/TiDBTestImages.java b/modules/tidb/src/test/java/org/testcontainers/TiDBTestImages.java
@@ -0,0 +1,8 @@
+package org.testcontainers;
+
+import org.testcontainers.utility.DockerImageName;
+
+public class TiDBTestImages {
+
+    public static final DockerImageName TIDB_IMAGE = DockerImageName.parse("pingcap/tidb:v6.1.0");
+}
diff --git a/modules/tidb/src/test/java/org/testcontainers/jdbc/tidb/TiDBJDBCDriverTest.java b/modules/tidb/src/test/java/org/testcontainers/jdbc/tidb/TiDBJDBCDriverTest.java
@@ -0,0 +1,19 @@
+package org.testcontainers.jdbc.tidb;
+
+import org.junit.runner.RunWith;
+import org.junit.runners.Parameterized;
+import org.testcontainers.jdbc.AbstractJDBCDriverTest;
+
+import java.util.Arrays;
+import java.util.EnumSet;
+
+@RunWith(Parameterized.class)
+public class TiDBJDBCDriverTest extends AbstractJDBCDriverTest {
+
+    @Parameterized.Parameters(name = "{index} - {0}")
+    public static Iterable<Object[]> data() {
+        return Arrays.asList(
+            new Object[][] { { "jdbc:tc:tidb://hostname/databasename", EnumSet.noneOf(Options.class) } }
+        );
+    }
+}
diff --git a/modules/tidb/src/test/java/org/testcontainers/junit/tidb/SimpleTiDBTest.java b/modules/tidb/src/test/java/org/testcontainers/junit/tidb/SimpleTiDBTest.java
@@ -0,0 +1,54 @@
+package org.testcontainers.junit.tidb;
+
+import org.hamcrest.Matchers;
+import org.junit.Assert;
+import org.junit.Test;
+import org.testcontainers.TiDBTestImages;
+import org.testcontainers.db.AbstractContainerDatabaseTest;
+import org.testcontainers.tidb.TiDBContainer;
+
+import java.sql.ResultSet;
+import java.sql.SQLException;
+
+public class SimpleTiDBTest extends AbstractContainerDatabaseTest {
+
+    @Test
+    public void testSimple() throws SQLException {
+        try (TiDBContainer tidb = new TiDBContainer(TiDBTestImages.TIDB_IMAGE)) {
+            tidb.start();
+
+            ResultSet resultSet = performQuery(tidb, "SELECT 1");
+
+            int resultSetInt = resultSet.getInt(1);
+            Assert.assertEquals("A basic SELECT query succeeds", 1, resultSetInt);
+        }
+    }
+
+    @Test
+    public void testExplicitInitScript() throws SQLException {
+        try (
+            TiDBContainer tidb = new TiDBContainer(TiDBTestImages.TIDB_IMAGE).withInitScript("somepath/init_tidb.sql")
+        ) { // TiDB is expected to be compatible with MySQL
+            tidb.start();
+
+            ResultSet resultSet = performQuery(tidb, "SELECT foo FROM bar");
+
+            String firstColumnValue = resultSet.getString(1);
+            Assert.assertEquals("Value from init script should equal real value", "hello world", firstColumnValue);
+        }
+    }
+
+    @Test
+    public void testWithAdditionalUrlParamInJdbcUrl() {
+        TiDBContainer tidb = new TiDBContainer(TiDBTestImages.TIDB_IMAGE).withUrlParam("sslmode", "disable");
+
+        try {
+            tidb.start();
+            String jdbcUrl = tidb.getJdbcUrl();
+            Assert.assertThat(jdbcUrl, Matchers.containsString("?"));
+            Assert.assertThat(jdbcUrl, Matchers.containsString("sslmode=disable"));
+        } finally {
+            tidb.stop();
+        }
+    }
+}
diff --git a/modules/tidb/src/test/resources/logback-test.xml b/modules/tidb/src/test/resources/logback-test.xml
@@ -0,0 +1,16 @@
+<configuration>
+
+    <appender name="STDOUT" class="ch.qos.logback.core.ConsoleAppender">
+        <!-- encoders are assigned the type
+             ch.qos.logback.classic.encoder.PatternLayoutEncoder by default -->
+        <encoder>
+            <pattern>%d{HH:mm:ss.SSS} %-5level %logger - %msg%n</pattern>
+        </encoder>
+    </appender>
+
+    <root level="INFO">
+        <appender-ref ref="STDOUT"/>
+    </root>
+
+    <logger name="org.testcontainers" level="DEBUG"/>
+</configuration>
diff --git a/modules/tidb/src/test/resources/somepath/init_tidb.sql b/modules/tidb/src/test/resources/somepath/init_tidb.sql
@@ -0,0 +1,16 @@
+CREATE TABLE bar (
+  foo VARCHAR(255)
+);
+
+SELECT "a /* string literal containing comment characters like -- here";
+SELECT "a 'quoting' \"scenario ` involving BEGIN keyword\" here";
+SELECT * from `bar`;
+
+-- What about a line comment containing imbalanced string delimiters? "
+
+/* or a block comment
+    containing imbalanced string delimiters?
+    ' "
+    */
+
+INSERT INTO bar (foo) /* ; */ VALUES ('hello world');
