瀏覽代碼

支持自定义脚本运行器参数.

nieqiurong 1 月之前
父節點
當前提交
36ea865c97

+ 85 - 6
mybatis-plus-extension/src/main/java/com/baomidou/mybatisplus/extension/ddl/DdlHelper.java

@@ -35,6 +35,7 @@ import java.time.format.DateTimeFormatter;
 import java.util.List;
 import java.util.Map;
 import java.util.Objects;
+import java.util.function.Consumer;
 
 /**
  * DDL 辅助类
@@ -50,7 +51,7 @@ public class DdlHelper {
      * 运行 SQL 脚本文件
      *
      * @param ddlGenerator DDL 生成器
-     * @param connection   数据库连接
+     * @param connection   数据库连接 (自行控制回收)
      * @param sqlFiles     SQL 文件列表
      * @param autoCommit   是否自动提交事务
      * @throws SQLException SQLException
@@ -62,19 +63,57 @@ public class DdlHelper {
     /**
      * 运行 SQL 脚本文件
      *
-     * @param ddlGenerator             DDL 生成器
-     * @param connection               数据库连接
-     * @param sqlFiles                 SQL 文件列表
-     * @param autoCommit               是否自动提交事务
+     * @param ddlGenerator          DDL 生成器
+     * @param connection            数据库连接 (自行控制回收)
+     * @param sqlFiles              SQL 文件列表
+     * @param autoCommit            是否自动提交事务
      * @param ddlScriptErrorHandler 错误处理器
      * @throws SQLException SQLException
      * @since 3.5.11
      */
-    public static void runScript(IDdlGenerator ddlGenerator, Connection connection, List<String> sqlFiles, boolean autoCommit, DdlScriptErrorHandler ddlScriptErrorHandler) throws SQLException {
+    public static void runScript(IDdlGenerator ddlGenerator, Connection connection, List<String> sqlFiles,
+                                 boolean autoCommit, DdlScriptErrorHandler ddlScriptErrorHandler) throws SQLException {
+        runScript(ddlGenerator, connection, sqlFiles, null, autoCommit, ddlScriptErrorHandler);
+    }
+
+    /**
+     * 运行 SQL 脚本文件
+     *
+     * @param ddlGenerator          DDL 生成器
+     * @param connection            数据库连接 (自行控制回收)
+     * @param sqlFiles              SQL 文件列表
+     * @param scriptRunnerConsumer  自定义 ScriptRunner 函数
+     * @param ddlScriptErrorHandler 错误处理器
+     * @throws SQLException SQLException
+     * @since 3.5.11
+     */
+    public static void runScript(IDdlGenerator ddlGenerator, Connection connection, List<String> sqlFiles,
+                                 Consumer<ScriptRunner> scriptRunnerConsumer, DdlScriptErrorHandler ddlScriptErrorHandler) throws SQLException {
+        runScript(ddlGenerator, connection, sqlFiles, scriptRunnerConsumer, false, ddlScriptErrorHandler);
+    }
+
+
+    /**
+     * 运行 SQL 脚本文件
+     *
+     * @param ddlGenerator          DDL 生成器
+     * @param connection            数据库连接 (自行控制回收)
+     * @param sqlFiles              SQL 文件列表
+     * @param scriptRunnerConsumer  自定义 ScriptRunner 函数
+     * @param autoCommit            是否自动提交事务
+     * @param ddlScriptErrorHandler 错误处理器
+     * @throws SQLException SQLException
+     * @since 3.5.11
+     */
+    public static void runScript(IDdlGenerator ddlGenerator, Connection connection, List<String> sqlFiles, Consumer<ScriptRunner> scriptRunnerConsumer,
+                                 boolean autoCommit, DdlScriptErrorHandler ddlScriptErrorHandler) throws SQLException {
         final String jdbcUrl = connection.getMetaData().getURL();
         final String schema = DdlHelper.getDatabase(jdbcUrl);
         SqlRunner sqlRunner = new SqlRunner(connection);
         ScriptRunner scriptRunner = getScriptRunner(connection, autoCommit);
+        if (scriptRunnerConsumer != null) {
+            scriptRunnerConsumer.accept(scriptRunner);
+        }
         if (null == ddlGenerator) {
             ddlGenerator = getDdlGenerator(jdbcUrl);
         }
@@ -128,7 +167,11 @@ public class DdlHelper {
      * @param dataSource   数据源
      * @param sqlFiles     SQL 文件列表
      * @param autoCommit   是否自动提交事务
+     * @see #runScript(IDdlGenerator, Connection, List, boolean)
+     * @see #runScript(IDdlGenerator, DataSource, List, boolean, DdlScriptErrorHandler)
+     * @deprecated 3.5.11 方法会吞掉所有异常,建议自行处理.
      */
+    @Deprecated
     public static void runScript(IDdlGenerator ddlGenerator, DataSource dataSource, List<String> sqlFiles, boolean autoCommit) {
         try (Connection connection = dataSource.getConnection()) {
             runScript(ddlGenerator, connection, sqlFiles, autoCommit);
@@ -155,6 +198,42 @@ public class DdlHelper {
         }
     }
 
+
+    /**
+     * 运行 SQL 脚本文件
+     *
+     * @param ddlGenerator          DDL 生成器
+     * @param dataSource            数据源
+     * @param sqlFiles              SQL 文件列表
+     * @param scriptRunnerConsumer  自定义 ScriptRunner 处理函数
+     * @param ddlScriptErrorHandler 错误处理器
+     * @since 3.5.11
+     */
+    public static void runScript(IDdlGenerator ddlGenerator,
+                                 DataSource dataSource, List<String> sqlFiles, Consumer<ScriptRunner> scriptRunnerConsumer,
+                                 DdlScriptErrorHandler ddlScriptErrorHandler) throws SQLException {
+        try (Connection connection = dataSource.getConnection()) {
+            runScript(ddlGenerator, connection, sqlFiles, scriptRunnerConsumer, false, ddlScriptErrorHandler);
+        }
+    }
+
+    /**
+     * 运行 SQL 脚本文件
+     *
+     * @param ddlGenerator          DDL 生成器
+     * @param dataSource            数据源
+     * @param sqlFiles              SQL 文件列表
+     * @param scriptRunnerConsumer  自定义 ScriptRunner 处理函数
+     * @param ddlScriptErrorHandler 错误处理器
+     * @since 3.5.11
+     */
+    public static void runScript(IDdlGenerator ddlGenerator, DataSource dataSource, List<String> sqlFiles,
+                                 Consumer<ScriptRunner> scriptRunnerConsumer, boolean autoCommit, DdlScriptErrorHandler ddlScriptErrorHandler) throws SQLException {
+        try (Connection connection = dataSource.getConnection()) {
+            runScript(ddlGenerator, connection, sqlFiles, scriptRunnerConsumer, autoCommit, ddlScriptErrorHandler);
+        }
+    }
+
     public static InputStream getInputStream(String path) throws Exception {
         return Resources.getResourceAsStream(path);
     }

+ 28 - 3
mybatis-plus-extension/src/main/java/com/baomidou/mybatisplus/extension/ddl/DdlScript.java

@@ -61,6 +61,13 @@ public class DdlScript {
      */
     private boolean autoCommit;
 
+    /**
+     * 自定义脚本运行器
+     *
+     * @since 3.5.11
+     */
+    private Consumer<ScriptRunner> scriptRunnerConsumer;
+
     /**
      * 非池化执行 (非自动提交)
      *
@@ -138,7 +145,24 @@ public class DdlScript {
      */
     @Deprecated
     public void run(List<String> sqlFiles, boolean autoCommit) {
-        DdlHelper.runScript(this.ddlGenerator, this.dataSource, sqlFiles, autoCommit);
+        try (Connection connection = this.dataSource.getConnection()) {
+            DdlHelper.runScript(this.ddlGenerator, connection, sqlFiles, this.scriptRunnerConsumer, autoCommit, DdlScriptErrorHandler.PrintlnLogErrorHandler.INSTANCE);
+        } catch (Exception e) {
+            // TODO 保持兼容,吞掉所有异常
+            LOG.error("Run script error: ", e);
+        }
+    }
+
+    /**
+     * 自定义 ScriptRunner
+     *
+     * @param scriptRunnerConsumer 处理函数
+     * @return this
+     * @since 3.5.11
+     */
+    public DdlScript scriptRunner(Consumer<ScriptRunner> scriptRunnerConsumer) {
+        this.scriptRunnerConsumer = scriptRunnerConsumer;
+        return this;
     }
 
     /**
@@ -186,11 +210,12 @@ public class DdlScript {
      */
     public void run(Connection connection, Reader reader, boolean autoCommit, String delimiter) {
         ScriptRunner scriptRunner = DdlHelper.getScriptRunner(connection, autoCommit);
-        // 设置自定义 SQL 分隔符,默认 ; 符号分割
+        if (scriptRunnerConsumer != null) {
+            scriptRunnerConsumer.accept(scriptRunner);
+        }
         if (StringUtils.isNotBlank(delimiter)) {
             scriptRunner.setDelimiter(delimiter);
         }
-        // 执行 SQL 脚本
         scriptRunner.runScript(reader);
     }
 

+ 8 - 6
mybatis-plus-extension/src/test/java/com/baomidou/mybatisplus/test/DdlScriptTest.java

@@ -12,17 +12,19 @@ public class DdlScriptTest {
 
     @Test
     void test() throws Exception {
-        var dddScript = new DdlScript(Driver.class.getName(),
+        var ddlScript = new DdlScript(Driver.class.getName(),
             "jdbc:h2:mem:test;MODE=mysql;DB_CLOSE_DELAY=-1;DB_CLOSE_ON_EXIT=FALSE",
-            "sa", "");
+            "sa", "", true).scriptRunner(scriptRunner -> {
+                scriptRunner.setLogWriter(null);
+        });
         LOGGER.info("--------------execute----------------");
-        dddScript.run( "select 1 from dual;", msg ->{});
+        ddlScript.run( "select 1 from dual;", msg ->{});
         LOGGER.info("--------------execute----------------");
-        dddScript.run( "select 2 from dual;", msg ->{});
+        ddlScript.run( "select 2 from dual;", msg ->{});
         LOGGER.info("--------------run----------------");
-        dddScript.run("select 1 from dual;");
+        ddlScript.run("select 1 from dual;");
         LOGGER.info("--------------run----------------");
-        dddScript.run("select 3 from dual#", "#");
+        ddlScript.run("select 3 from dual#", "#");
     }
 
 }