Add JDBC custom credential support for postgres (testcontainers#1293)

Alexander198961 · Mar 17, 2019 · 1fe2804 · 1fe2804
1 parent 91ca680
commit 1fe2804
Show file tree

Hide file tree

Showing 6 changed files with 70 additions and 43 deletions.
diff --git a/modules/jdbc-test/src/test/java/org/testcontainers/jdbc/JDBCDriverTest.java b/modules/jdbc-test/src/test/java/org/testcontainers/jdbc/JDBCDriverTest.java
@@ -50,9 +50,9 @@ public static Iterable<Object[]> data() {
                 {"jdbc:tc:mysql:5.5.43://hostname/databasename?TC_INITSCRIPT=somepath/init_unicode_mysql.sql&useUnicode=yes&characterEncoding=utf8", EnumSet.of(Options.CharacterSet)},
                 {"jdbc:tc:mysql:5.5.43://hostname/databasename", EnumSet.noneOf(Options.class)},
                 {"jdbc:tc:mysql:5.5.43://hostname/databasename?useSSL=false", EnumSet.noneOf(Options.class)},
-                {"jdbc:tc:postgresql:9.6.8://hostname/databasename", EnumSet.noneOf(Options.class)},
-                {"jdbc:tc:postgis://hostname/databasename", EnumSet.noneOf(Options.class)},
-                {"jdbc:tc:postgis:9.6://hostname/databasename", EnumSet.noneOf(Options.class)},
+                {"jdbc:tc:postgresql:9.6.8://hostname/databasename?user=someuser&password=somepwd", EnumSet.of(Options.JDBCParams)},
+                {"jdbc:tc:postgis://hostname/databasename?user=someuser&password=somepwd", EnumSet.of(Options.JDBCParams)},
+                {"jdbc:tc:postgis:9.6://hostname/databasename?user=someuser&password=somepwd", EnumSet.of(Options.JDBCParams)},
                 {"jdbc:tc:mysql:5.6://hostname/databasename?TC_MY_CNF=somepath/mysql_conf_override", EnumSet.of(Options.CustomIniFile)},
                 {"jdbc:tc:mariadb://hostname/databasename", EnumSet.noneOf(Options.class)},
                 {"jdbc:tc:mariadb://hostname/databasename?user=someuser&TC_INITSCRIPT=somepath/init_mariadb.sql", EnumSet.of(Options.ScriptedSchema, Options.JDBCParams)},
@@ -135,16 +135,27 @@ private void performTestForScriptedSchema(String jdbcUrl) throws SQLException {
 
     private void performTestForJDBCParamUsage(String jdbcUrl) throws SQLException {
         try (HikariDataSource dataSource = getDataSource(jdbcUrl, 1)) {
-            boolean result = new QueryRunner(dataSource).query("select CURRENT_USER()", rs -> {
+            boolean result = new QueryRunner(dataSource).query("select CURRENT_USER", rs -> {
                 rs.next();
                 String resultUser = rs.getString(1);
-                assertEquals("User from query param is created.", "someuser@%", resultUser);
+                // Not all databases (eg. Postgres) return @% at the end of user name. We just need to make sure the user name matches.
+                if (resultUser.endsWith("@%")) {
+                    resultUser = resultUser.substring(0, resultUser.length() - 2);
+                }
+                assertEquals("User from query param is created.", "someuser", resultUser);
                 return true;
             });
 
             assertTrue("The database returned a record as expected", result);
 
-            result = new QueryRunner(dataSource).query("SELECT DATABASE()", rs -> {
+            String databaseQuery = "SELECT DATABASE()";
+            // Postgres does not have Database() as a function
+            String databaseType = ConnectionUrl.newInstance(jdbcUrl).getDatabaseType();
+            if (databaseType.equalsIgnoreCase("postgresql") || databaseType.equalsIgnoreCase("postgis")) {
+                databaseQuery = "SELECT CURRENT_DATABASE()";
+            }
+
+            result = new QueryRunner(dataSource).query(databaseQuery, rs -> {
                 rs.next();
                 String resultDB = rs.getString(1);
                 assertEquals("Database name from URL String is used.", "databasename", resultDB);

diff --git a/modules/jdbc/src/main/java/org/testcontainers/containers/JdbcDatabaseContainerProvider.java b/modules/jdbc/src/main/java/org/testcontainers/containers/JdbcDatabaseContainerProvider.java
@@ -4,6 +4,8 @@
 
 import lombok.extern.slf4j.Slf4j;
 
+import java.util.Objects;
+
 /**
  * Base class for classes that can provide a JDBC container.
  */
@@ -49,4 +51,24 @@ public JdbcDatabaseContainer newInstance(ConnectionUrl url) {
             return newInstance();
         }
     }
+
+    protected JdbcDatabaseContainer newInstanceFromConnectionUrl(ConnectionUrl connectionUrl, final String userParamName, final String pwdParamName) {
+        Objects.requireNonNull(connectionUrl, "Connection URL cannot be null");
+
+        final String databaseName = connectionUrl.getDatabaseName().orElse("test");
+        final String user = connectionUrl.getQueryParameters().getOrDefault(userParamName, "test");
+        final String password = connectionUrl.getQueryParameters().getOrDefault(pwdParamName, "test");
+
+        final JdbcDatabaseContainer instance;
+        if (connectionUrl.getImageTag().isPresent()) {
+            instance = newInstance(connectionUrl.getImageTag().get());
+        } else {
+            instance = newInstance();
+        }
+
+        return instance
+            .withDatabaseName(databaseName)
+            .withUsername(user)
+            .withPassword(password);
+    }
 }
diff --git a/modules/mariadb/src/main/java/org/testcontainers/containers/MariaDBContainerProvider.java b/modules/mariadb/src/main/java/org/testcontainers/containers/MariaDBContainerProvider.java
@@ -8,7 +8,7 @@
  * Factory for MariaDB org.testcontainers.containers.
  */
 public class MariaDBContainerProvider extends JdbcDatabaseContainerProvider {
-    
+
     private static final String USER_PARAM = "user";
 
     private static final String PASSWORD_PARAM = "password";
@@ -27,26 +27,10 @@ public JdbcDatabaseContainer newInstance() {
     public JdbcDatabaseContainer newInstance(String tag) {
         return new MariaDBContainer(MariaDBContainer.IMAGE + ":" + tag);
     }
-    
+
     @Override
     public JdbcDatabaseContainer newInstance(ConnectionUrl connectionUrl) {
-        Objects.requireNonNull(connectionUrl, "Connection URL cannot be null");
-
-        final String databaseName = connectionUrl.getDatabaseName().orElse("test");
-        final String user = connectionUrl.getQueryParameters().getOrDefault(USER_PARAM, "test");
-        final String password = connectionUrl.getQueryParameters().getOrDefault(PASSWORD_PARAM, "test");
-
-        final JdbcDatabaseContainer instance;
-        if (connectionUrl.getImageTag().isPresent()) {
-            instance = newInstance(connectionUrl.getImageTag().get());
-        } else {
-            instance = newInstance();
-        }
-
-        return instance
-            .withDatabaseName(databaseName)
-            .withUsername(user)
-            .withPassword(password);
+       return newInstanceFromConnectionUrl(connectionUrl, USER_PARAM, PASSWORD_PARAM);
     }
-    
+
 }
diff --git a/modules/mysql/src/main/java/org/testcontainers/containers/MySQLContainerProvider.java b/modules/mysql/src/main/java/org/testcontainers/containers/MySQLContainerProvider.java
@@ -34,23 +34,7 @@ public JdbcDatabaseContainer newInstance(String tag) {
 
     @Override
     public JdbcDatabaseContainer newInstance(ConnectionUrl connectionUrl) {
-        Objects.requireNonNull(connectionUrl, "Connection URL cannot be null");
-
-        final String databaseName = connectionUrl.getDatabaseName().orElse("test");
-        final String user = connectionUrl.getQueryParameters().getOrDefault(USER_PARAM, "test");
-        final String password = connectionUrl.getQueryParameters().getOrDefault(PASSWORD_PARAM, "test");
-
-        final JdbcDatabaseContainer instance;
-        if (connectionUrl.getImageTag().isPresent()) {
-            instance = newInstance(connectionUrl.getImageTag().get());
-        } else {
-            instance = newInstance();
-        }
-
-        return instance
-            .withDatabaseName(databaseName)
-            .withUsername(user)
-            .withPassword(password);
+        return newInstanceFromConnectionUrl(connectionUrl, USER_PARAM, PASSWORD_PARAM);
     }
 
 }
diff --git a/modules/postgresql/src/main/java/org/testcontainers/containers/PostgisContainerProvider.java b/modules/postgresql/src/main/java/org/testcontainers/containers/PostgisContainerProvider.java
@@ -1,5 +1,9 @@
 package org.testcontainers.containers;
 
+import org.testcontainers.jdbc.ConnectionUrl;
+
+import java.util.Objects;
+
 /**
  * Factory for PostGIS containers, which are a special flavour of PostgreSQL.
  */
@@ -8,6 +12,9 @@ public class PostgisContainerProvider extends JdbcDatabaseContainerProvider {
     private static final String NAME = "postgis";
     private static final String DEFAULT_TAG = "10";
     private static final String DEFAULT_IMAGE = "mdillon/postgis";
+    public static final String USER_PARAM = "user";
+    public static final String PASSWORD_PARAM = "password";
+
 
     @Override
     public boolean supports(String databaseType) {
@@ -23,4 +30,9 @@ public JdbcDatabaseContainer newInstance() {
     public JdbcDatabaseContainer newInstance(String tag) {
         return new PostgreSQLContainer(DEFAULT_IMAGE + ":" + tag);
     }
+
+    @Override
+    public JdbcDatabaseContainer newInstance(ConnectionUrl connectionUrl) {
+        return newInstanceFromConnectionUrl(connectionUrl, USER_PARAM, PASSWORD_PARAM);
+    }
 }
diff --git a/...s/postgresql/src/main/java/org/testcontainers/containers/PostgreSQLContainerProvider.java b/...s/postgresql/src/main/java/org/testcontainers/containers/PostgreSQLContainerProvider.java
@@ -1,9 +1,17 @@
 package org.testcontainers.containers;
 
+import org.testcontainers.jdbc.ConnectionUrl;
+
+import java.util.Objects;
+
 /**
  * Factory for PostgreSQL containers.
  */
 public class PostgreSQLContainerProvider extends JdbcDatabaseContainerProvider {
+
+    public static final String USER_PARAM = "user";
+    public static final String PASSWORD_PARAM = "password";
+
     @Override
     public boolean supports(String databaseType) {
         return databaseType.equals(PostgreSQLContainer.NAME);
@@ -18,4 +26,10 @@ public JdbcDatabaseContainer newInstance() {
     public JdbcDatabaseContainer newInstance(String tag) {
         return new PostgreSQLContainer(PostgreSQLContainer.IMAGE + ":" + tag);
     }
+
+    @Override
+    public JdbcDatabaseContainer newInstance(ConnectionUrl connectionUrl) {
+        return newInstanceFromConnectionUrl(connectionUrl, USER_PARAM, PASSWORD_PARAM);
+    }
+
 }