Merge branch '7.6'

gchq · Oct 31, 2024 · 13e3f3a · 13e3f3a
2 parents 03a6827 + fca066b
commit 13e3f3a
Show file tree

Hide file tree

Showing 4 changed files with 471 additions and 24 deletions.
diff --git a/.idea/codeStyles/Project.xml b/.idea/codeStyles/Project.xml
diff --git a/stroom-test-common/build.gradle b/stroom-test-common/build.gradle
@@ -29,6 +29,7 @@ dependencies {
     implementation libs.vavr
     implementation libs.wix_embedded_mysql
 
+    runtimeOnly libs.mysql_connector_java
 
     // The following logging libs are needed when running junits outside dropwizard
     runtimeOnly libs.jcl_over_slf4j

diff --git a/stroom-test-common/src/main/java/stroom/test/common/util/db/DbTestUtil.java b/stroom-test-common/src/main/java/stroom/test/common/util/db/DbTestUtil.java
@@ -10,6 +10,7 @@
 import stroom.util.ConsoleColour;
 import stroom.util.NullSafe;
 import stroom.util.db.ForceLegacyMigration;
+import stroom.util.exception.ThrowingConsumer;
 import stroom.util.io.FileUtil;
 import stroom.util.logging.LambdaLogger;
 import stroom.util.logging.LambdaLoggerFactory;
@@ -47,6 +48,7 @@
 import java.util.concurrent.ConcurrentMap;
 import java.util.concurrent.ConcurrentSkipListSet;
 import java.util.concurrent.TimeUnit;
+import java.util.function.Consumer;
 import java.util.function.Function;
 import java.util.function.Predicate;
 import java.util.function.Supplier;
@@ -226,19 +228,42 @@ public static void dropUnusedTestDatabases() {
         }
     }
 
-    public static void dropAllTestDatabases() {
+    public static void doWithRootConnection(final Consumer<Connection> connectionConsumer) {
+        if (connectionConsumer != null) {
+            getWithRootConnection(connection -> {
+                connectionConsumer.accept(connection);
+                // Result ignored
+                return null;
+            });
+        }
+    }
+
+    public static <T> T getWithRootConnection(final Function<Connection, T> conectionFunction) {
+        Objects.requireNonNull(conectionFunction);
         final ConnectionConfig connectionConfig = createConnectionConfig(new CommonDbConfig());
         final ConnectionConfig rootConnectionConfig = createRootConnectionConfig(connectionConfig);
 
-        // Create new db.
-        final Properties connectionProps = new Properties();
-        connectionProps.put("user", rootConnectionConfig.getUser());
-        connectionProps.put("password", rootConnectionConfig.getPassword());
+        ensureJdbcDriver(connectionConfig);
+        try (final Connection connection = DriverManager.getConnection(rootConnectionConfig.getUrl(),
+                rootConnectionConfig.getUser(),
+                rootConnectionConfig.getPassword())) {
+            return conectionFunction.apply(connection);
+        } catch (final SQLException e) {
+            throw new RuntimeException(e.getMessage(), e);
+        }
+    }
 
-        final Predicate<String> dbNameMatchPredicate = DB_NAME_PATTERN.asMatchPredicate();
+    public static void ensureJdbcDriver(final ConnectionConfig connectionConfig) {
+        try {
+            Class.forName(connectionConfig.getClassName());
+        } catch (ClassNotFoundException e) {
+            throw new RuntimeException(e);
+        }
+    }
 
-        try (final Connection connection = DriverManager.getConnection(rootConnectionConfig.getUrl(),
-                connectionProps)) {
+    public static void dropAllTestDatabases() {
+        final Predicate<String> dbNameMatchPredicate = DB_NAME_PATTERN.asMatchPredicate();
+        doWithRootConnection(ThrowingConsumer.unchecked(connection -> {
             try (final Statement statement = connection.createStatement()) {
                 final ResultSet resultSet = statement.executeQuery("SHOW DATABASES;");
                 final List<String> dbNames = new ArrayList<>();
@@ -254,9 +279,17 @@ public static void dropAllTestDatabases() {
                     statement.executeUpdate("DROP DATABASE " + dbName + ";");
                 }
             }
-        } catch (final SQLException e) {
-            throw new RuntimeException(e.getMessage(), e);
-        }
+        }));
+    }
+
+    public static void dropDatabase(final String dbName) {
+        doWithRootConnection(ThrowingConsumer.unchecked(rootConn -> {
+            try (final Statement statement = rootConn.createStatement()) {
+                statement.executeUpdate("DROP DATABASE `" + dbName + "`;");
+            }
+        }));
+
+
     }
 
 //    private static String getThreadLocalDbName(final boolean isSharedDatabase) {
@@ -415,18 +448,9 @@ public static DataSource createTestDataSource(final AbstractDbConfig dbConfig,
 
             try (final Connection connection = DriverManager.getConnection(rootConnectionConfig.getUrl(),
                     connectionProps)) {
-                try (final Statement statement = connection.createStatement()) {
-
-                    statement.executeUpdate("CREATE DATABASE `" + dbName +
-                            "` CHARACTER SET = utf8mb4 COLLATE = utf8mb4_unicode_ci;");
-
-                    statement.executeUpdate("CREATE USER IF NOT EXISTS '" +
-                            connectionConfig.getUser() + "'@'%' IDENTIFIED BY '" +
-                            connectionConfig.getPassword() + "';");
-
-                    statement.executeUpdate("GRANT ALL PRIVILEGES ON *.* TO '" +
-                            connectionConfig.getUser() + "'@'%' WITH GRANT OPTION;");
-                }
+                final String username = connectionConfig.getUser();
+                final String password = connectionConfig.getPassword();
+                createStroomDatabaseAndUser(connection, dbName, username, password);
             } catch (final SQLException e) {
                 throw new RuntimeException(e.getMessage(), e);
             }
@@ -471,6 +495,26 @@ public static DataSource createTestDataSource(final AbstractDbConfig dbConfig,
         return dataSource;
     }
 
+    public static void createStroomDatabaseAndUser(final Connection connection,
+                                                   final String dbName,
+                                                   final String username,
+                                                   final String password) throws SQLException {
+        try (final Statement statement = connection.createStatement()) {
+            LOGGER.debug("Creating database '{}'", dbName);
+            statement.executeUpdate("CREATE DATABASE `" + dbName +
+                    "` CHARACTER SET = utf8mb4 COLLATE = utf8mb4_unicode_ci;");
+
+            LOGGER.debug("Creating DB user '{}'", username);
+            statement.executeUpdate("CREATE USER IF NOT EXISTS '" +
+                    username + "'@'%' IDENTIFIED BY '" +
+                    password + "';");
+
+            LOGGER.debug("Granting privileges to user '{}'", username);
+            statement.executeUpdate("GRANT ALL PRIVILEGES ON *.* TO '" +
+                    username + "'@'%' WITH GRANT OPTION;");
+        }
+    }
+
     private static DataSource createDataSource(final String name,
                                                final boolean unique,
                                                final AbstractDbConfig testDbConfig) {