优化代码生成器

mybatis-plus-extension/build.gradle

@@ -1,3 +1,4 @@
+apply plugin: 'kotlin'
 
 dependencies {
     api project(":mybatis-plus-core")

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

@@ -50,52 +50,65 @@ public class DdlHelper {
      * 允许 SQL 脚本文件
      *
      * @param ddlGenerator DDL 生成器
-     * @param dataSource   数据源
+     * @param connection   数据库连接
      * @param sqlFiles     SQL 文件列表
      * @param autoCommit   字段提交事务
      * @throws Exception
      */
-    public static void runScript(IDdlGenerator ddlGenerator, DataSource dataSource, List<String> sqlFiles, boolean autoCommit) {
-        try (Connection connection = dataSource.getConnection()) {
-            // 执行自定义 DDL 信息
-            final String jdbcUrl = connection.getMetaData().getURL();
-            final String schema = DdlHelper.getDatabase(jdbcUrl);
-            SqlRunner sqlRunner = new SqlRunner(connection);
-            ScriptRunner scriptRunner = getScriptRunner(connection, autoCommit);
-            if (null == ddlGenerator) {
-                ddlGenerator = getDdlGenerator(jdbcUrl);
-            }
-            if (!ddlGenerator.existTable(schema, sql -> {
-                try {
-                    Map<String, Object> resultMap = sqlRunner.selectOne(sql);
-                    if (null != resultMap && !StringPool.ZERO.equals(String.valueOf(resultMap.get(StringPool.NUM)))) {
-                        return true;
-                    }
-                } catch (SQLException e) {
-                    log.error("run script sql:{} , error: {}", sql, e.getMessage());
+    public static void runScript(IDdlGenerator ddlGenerator, Connection connection, List<String> sqlFiles, boolean autoCommit) throws SQLException {
+        // 执行自定义 DDL 信息
+        final String jdbcUrl = connection.getMetaData().getURL();
+        final String schema = DdlHelper.getDatabase(jdbcUrl);
+        SqlRunner sqlRunner = new SqlRunner(connection);
+        ScriptRunner scriptRunner = getScriptRunner(connection, autoCommit);
+        if (null == ddlGenerator) {
+            ddlGenerator = getDdlGenerator(jdbcUrl);
+        }
+        if (!ddlGenerator.existTable(schema, sql -> {
+            try {
+                Map<String, Object> resultMap = sqlRunner.selectOne(sql);
+                if (null != resultMap && !StringPool.ZERO.equals(String.valueOf(resultMap.get(StringPool.NUM)))) {
+                    return true;
                 }
-                return false;
-            })) {
-                scriptRunner.runScript(new StringReader(ddlGenerator.createDdlHistory()));
+            } catch (SQLException e) {
+                log.error("run script sql:{} , error: {}", sql, e.getMessage());
             }
-            // 执行 SQL 脚本
-            for (String sqlFile : sqlFiles) {
-                try {
-                    List<Map<String, Object>> objectMap = sqlRunner.selectAll(ddlGenerator.selectDdlHistory(sqlFile, StringPool.SQL));
-                    if (null == objectMap || objectMap.isEmpty()) {
-                        log.debug("run script file: {}", sqlFile);
-                        File file = new File(sqlFile);
-                        if (file.exists()) {
-                            scriptRunner.runScript(new FileReader(file));
-                        } else {
-                            scriptRunner.runScript(new InputStreamReader(getInputStream(sqlFile)));
-                        }
-                        sqlRunner.insert(ddlGenerator.insertDdlHistory(sqlFile, StringPool.SQL, getNowTime()));
+            return false;
+        })) {
+            scriptRunner.runScript(new StringReader(ddlGenerator.createDdlHistory()));
+        }
+        // 执行 SQL 脚本
+        for (String sqlFile : sqlFiles) {
+            try {
+                List<Map<String, Object>> objectMap = sqlRunner.selectAll(ddlGenerator.selectDdlHistory(sqlFile, StringPool.SQL));
+                if (null == objectMap || objectMap.isEmpty()) {
+                    log.debug("run script file: {}", sqlFile);
+                    File file = new File(sqlFile);
+                    if (file.exists()) {
+                        scriptRunner.runScript(new FileReader(file));
+                    } else {
+                        scriptRunner.runScript(new InputStreamReader(getInputStream(sqlFile)));
                     }
-                } catch (Exception e) {
-                    log.error("run script sql:{} , error: {} , Please check if the table `ddl_history` exists", sqlFile, e.getMessage());
+                    sqlRunner.insert(ddlGenerator.insertDdlHistory(sqlFile, StringPool.SQL, getNowTime()));
                 }
+            } catch (Exception e) {
+                log.error("run script sql:{} , error: {} , Please check if the table `ddl_history` exists", sqlFile, e.getMessage());
             }
+        }
+    }
+
+    /**
+     * 允许 SQL 脚本文件
+     *
+     * @param ddlGenerator DDL 生成器
+     * @param dataSource   数据源
+     * @param sqlFiles     SQL 文件列表
+     * @param autoCommit   字段提交事务
+     * @throws Exception
+     */
+    public static void runScript(IDdlGenerator ddlGenerator, DataSource dataSource, List<String> sqlFiles, boolean autoCommit) {
+        try (Connection connection = dataSource.getConnection()) {
+            runScript(ddlGenerator, connection, sqlFiles, autoCommit);
         } catch (Exception e) {
             log.error("run script error: {}", e.getMessage());
         }

mybatis-plus-generator/src/main/java/com/baomidou/mybatisplus/generator/jdbc/DatabaseMetaDataWrapper.java

@@ -17,7 +17,6 @@ package com.baomidou.mybatisplus.generator.jdbc;
 
 import com.baomidou.mybatisplus.core.toolkit.StringPool;
 import com.baomidou.mybatisplus.core.toolkit.StringUtils;
-import com.baomidou.mybatisplus.generator.config.DataSourceConfig;
 import org.apache.ibatis.type.JdbcType;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
@@ -38,6 +37,8 @@ public class DatabaseMetaDataWrapper {
 
     private static final Logger logger = LoggerFactory.getLogger(DatabaseMetaDataWrapper.class);
 
+    private Connection connection;
+
     private final DatabaseMetaData databaseMetaData;
 
     //TODO 暂时只支持一种
@@ -46,17 +47,34 @@ public class DatabaseMetaDataWrapper {
     //TODO 暂时只支持一种
     private final String schema;
 
-    public DatabaseMetaDataWrapper(DataSourceConfig dataSourceConfig) {
+    public DatabaseMetaDataWrapper(Connection connection, String schemaName) {
         try {
-            Connection connection = dataSourceConfig.getConn();
+            if (null == connection) {
+                throw new RuntimeException("数据库连接不能为空");
+            }
+            this.connection = connection;
             this.databaseMetaData = connection.getMetaData();
             this.catalog = connection.getCatalog();
-            this.schema = dataSourceConfig.getSchemaName();
+            this.schema = schemaName;
         } catch (SQLException e) {
             throw new RuntimeException("获取元数据错误:", e);
         }
     }
 
+    public Connection getConnection() {
+        return connection;
+    }
+
+    public void closeConnection() {
+        Optional.ofNullable(connection).ifPresent((con) -> {
+            try {
+                con.close();
+            } catch (SQLException sqlException) {
+                sqlException.printStackTrace();
+            }
+        });
+    }
+
     public Map<String, Column> getColumnsInfo(String tableNamePattern, boolean queryPrimaryKey) {
         return getColumnsInfo(this.catalog, this.schema, tableNamePattern, queryPrimaryKey);
     }

mybatis-plus-generator/src/main/java/com/baomidou/mybatisplus/generator/query/AbstractDatabaseQuery.java

@@ -50,17 +50,12 @@ public abstract class AbstractDatabaseQuery implements IDatabaseQuery {
 
     protected final GlobalConfig globalConfig;
 
-    protected final DbQueryDecorator dbQuery;
-
-    protected final DatabaseMetaDataWrapper databaseMetaDataWrapper;
 
     public AbstractDatabaseQuery(@NotNull ConfigBuilder configBuilder) {
         this.configBuilder = configBuilder;
         this.dataSourceConfig = configBuilder.getDataSourceConfig();
         this.strategyConfig = configBuilder.getStrategyConfig();
-        this.dbQuery = new DbQueryDecorator(dataSourceConfig, strategyConfig);
         this.globalConfig = configBuilder.getGlobalConfig();
-        this.databaseMetaDataWrapper = new DatabaseMetaDataWrapper(dataSourceConfig);
     }
 
     @NotNull

mybatis-plus-generator/src/main/java/com/baomidou/mybatisplus/generator/query/DefaultQuery.java

@@ -49,39 +49,46 @@ import java.util.Map;
 public class DefaultQuery extends AbstractDatabaseQuery {
 
     private final TypeRegistry typeRegistry;
+    protected final DatabaseMetaDataWrapper databaseMetaDataWrapper;
 
     public DefaultQuery(@NotNull ConfigBuilder configBuilder) {
         super(configBuilder);
         typeRegistry = new TypeRegistry(configBuilder.getGlobalConfig());
+        this.databaseMetaDataWrapper = new DatabaseMetaDataWrapper(dataSourceConfig.getConn(), dataSourceConfig.getSchemaName());
     }
 
     @Override
     public @NotNull List<TableInfo> queryTables() {
-        boolean isInclude = strategyConfig.getInclude().size() > 0;
-        boolean isExclude = strategyConfig.getExclude().size() > 0;
-        //所有的表信息
-        List<TableInfo> tableList = new ArrayList<>();
-        List<DatabaseMetaDataWrapper.Table> tables = getTables();
-        //需要反向生成或排除的表信息
-        List<TableInfo> includeTableList = new ArrayList<>();
-        List<TableInfo> excludeTableList = new ArrayList<>();
-        tables.forEach(table -> {
-            String tableName = table.getName();
-            if (StringUtils.isNotBlank(tableName)) {
-                TableInfo tableInfo = new TableInfo(this.configBuilder, tableName);
-                tableInfo.setComment(table.getRemarks());
-                if (isInclude && strategyConfig.matchIncludeTable(tableName)) {
-                    includeTableList.add(tableInfo);
-                } else if (isExclude && strategyConfig.matchExcludeTable(tableName)) {
-                    excludeTableList.add(tableInfo);
+        try {
+            boolean isInclude = strategyConfig.getInclude().size() > 0;
+            boolean isExclude = strategyConfig.getExclude().size() > 0;
+            //所有的表信息
+            List<TableInfo> tableList = new ArrayList<>();
+            List<DatabaseMetaDataWrapper.Table> tables = this.getTables();
+            //需要反向生成或排除的表信息
+            List<TableInfo> includeTableList = new ArrayList<>();
+            List<TableInfo> excludeTableList = new ArrayList<>();
+            tables.forEach(table -> {
+                String tableName = table.getName();
+                if (StringUtils.isNotBlank(tableName)) {
+                    TableInfo tableInfo = new TableInfo(this.configBuilder, tableName);
+                    tableInfo.setComment(table.getRemarks());
+                    if (isInclude && strategyConfig.matchIncludeTable(tableName)) {
+                        includeTableList.add(tableInfo);
+                    } else if (isExclude && strategyConfig.matchExcludeTable(tableName)) {
+                        excludeTableList.add(tableInfo);
+                    }
+                    tableList.add(tableInfo);
                 }
-                tableList.add(tableInfo);
-            }
-        });
-        filter(tableList, includeTableList, excludeTableList);
-        // 性能优化，只处理需执行表字段 https://github.com/baomidou/mybatis-plus/issues/219
-        tableList.forEach(this::convertTableFields);
-        return tableList;
+            });
+            filter(tableList, includeTableList, excludeTableList);
+            // 性能优化，只处理需执行表字段 https://github.com/baomidou/mybatis-plus/issues/219
+            tableList.forEach(this::convertTableFields);
+            return tableList;
+        } finally {
+            // 数据库操作完成,释放连接对象
+            databaseMetaDataWrapper.closeConnection();
+        }
     }
 
     protected List<DatabaseMetaDataWrapper.Table> getTables() {

mybatis-plus-generator/src/main/java/com/baomidou/mybatisplus/generator/query/SQLQuery.java

@@ -23,6 +23,7 @@ import com.baomidou.mybatisplus.generator.config.builder.ConfigBuilder;
 import com.baomidou.mybatisplus.generator.config.builder.Entity;
 import com.baomidou.mybatisplus.generator.config.po.TableField;
 import com.baomidou.mybatisplus.generator.config.po.TableInfo;
+import com.baomidou.mybatisplus.generator.config.querys.DbQueryDecorator;
 import com.baomidou.mybatisplus.generator.config.querys.H2Query;
 import com.baomidou.mybatisplus.generator.config.rules.IColumnType;
 import com.baomidou.mybatisplus.generator.jdbc.DatabaseMetaDataWrapper;
@@ -41,8 +42,13 @@ import java.util.*;
  */
 public class SQLQuery extends AbstractDatabaseQuery {
 
+    protected final DbQueryDecorator dbQuery;
+    protected final DatabaseMetaDataWrapper databaseMetaDataWrapper;
+
     public SQLQuery(@NotNull ConfigBuilder configBuilder) {
         super(configBuilder);
+        this.dbQuery = new DbQueryDecorator(dataSourceConfig, strategyConfig);
+        this.databaseMetaDataWrapper = new DatabaseMetaDataWrapper(dbQuery.getConnection(), dataSourceConfig.getSchemaName());
     }
 
     @NotNull

mybatis-plus-generator/src/test/java/com/baomidou/mybatisplus/generator/jdbc/DatabaseMetaDataWrapperTest.java

@@ -13,7 +13,7 @@ public class DatabaseMetaDataWrapperTest {
     @Test
     void test() throws SQLException {
         DataSourceConfig dataSourceConfig = new DataSourceConfig.Builder("jdbc:h2:mem:test;MODE=mysql;DB_CLOSE_DELAY=-1;DB_CLOSE_ON_EXIT=FALSE", "sa", "").build();
-        DatabaseMetaDataWrapper databaseMetaDataWrapper = new DatabaseMetaDataWrapper(dataSourceConfig);
+        DatabaseMetaDataWrapper databaseMetaDataWrapper = new DatabaseMetaDataWrapper(dataSourceConfig.getConn(), dataSourceConfig.getSchemaName());
         Map<String, DatabaseMetaDataWrapper.Column> columnsInfo = databaseMetaDataWrapper.getColumnsInfo(null, null, "USERS",true);
         Assertions.assertNotNull(columnsInfo);
         DatabaseMetaDataWrapper.Column name = columnsInfo.get("name");