cn.mybatis.mp.generator.database.meta.TableMetaDataQuery Maven / Gradle / Ivy

Go to download
package cn.mybatis.mp.generator.database.meta;

import cn.mybatis.mp.generator.config.GeneratorConfig;
import db.sql.api.impl.tookit.Objects;
import lombok.extern.slf4j.Slf4j;
import org.apache.ibatis.type.JdbcType;

import java.sql.Connection;
import java.sql.DatabaseMetaData;
import java.sql.ResultSet;
import java.sql.SQLException;
import java.util.ArrayList;
import java.util.HashSet;
import java.util.List;
import java.util.Set;
import java.util.stream.Collectors;

@Slf4j
public class TableMetaDataQuery {

    protected final DatabaseMetaData metaData;

    protected final GeneratorConfig generatorConfig;

    protected final String connDatabaseName;

    protected final String connSchema;

    public TableMetaDataQuery(GeneratorConfig generatorConfig, Connection connection) throws SQLException {
        this.metaData = connection.getMetaData();
        this.connSchema = connection.getSchema();
        this.generatorConfig = generatorConfig;
        this.connDatabaseName = connection.getCatalog();
    }

    private boolean isTableNeedInclude(String tableName) {
        if (generatorConfig.getTableConfig().getExcludeTables().contains(tableName)) {
            return false;
        }

        if (!generatorConfig.getTableConfig().getIncludeTables().isEmpty() && !generatorConfig.getTableConfig().getIncludeTables().contains(tableName)) {
            return false;
        }
        return true;
    }

    private List getTableInfoList(String databaseName, String schema, String tableNamePattern, String[] types) {
        List tables = new ArrayList<>();
        try (ResultSet resultSet = metaData.getTables(databaseName, schema, tableNamePattern, types)) {
            TableInfo tableInfo;
            while (resultSet.next()) {
                String TABLE_NAME = resultSet.getString("TABLE_NAME");
                String tableName = TABLE_NAME;

                if (tableName.toUpperCase().equals(tableName)) {
                    tableName = tableName.toLowerCase();
                }

                if (!isTableNeedInclude(TABLE_NAME) && !isTableNeedInclude(tableName)) {
                    continue;
                }

                tableInfo = new TableInfo();
                tableInfo.setName(tableName);
                tableInfo.setRemarks(resultSet.getString("REMARKS"));
                tableInfo.setTableType(resultSet.getString("TABLE_TYPE"));
                tableInfo.setSchema(resultSet.getString("TABLE_SCHEM"));
                tableInfo.setCatalog(resultSet.getString("TABLE_CAT"));

                tableInfo.setColumnInfoList(getColumnInfo(tableInfo, TABLE_NAME));
                List idColumnInfoList = tableInfo.getColumnInfoList().stream().filter(item -> item.isPrimaryKey()).collect(Collectors.toList());
                if (!idColumnInfoList.isEmpty() && idColumnInfoList.size() == 1) {
                    tableInfo.setIdColumnInfo(idColumnInfoList.get(0));
                }
                tableInfo.setIdColumnInfoList(idColumnInfoList);
                tables.add(tableInfo);
            }
        } catch (SQLException e) {
            throw new RuntimeException("读取数据库表信息出现错误", e);
        }

        return tables;
    }

    public List getTableInfoList(boolean includeTable, boolean includeView) {
        String[] types;
        if (includeTable && includeView) {
            types = new String[]{"TABLE", "VIEW"};
        } else {
            if (includeTable) {
                types = new String[]{"TABLE"};
            } else if (includeView) {
                types = new String[]{"VIEW"};
            } else {
                throw new RuntimeException("includeTable or includeView must be true");
            }
        }

        String databaseName = generatorConfig.getDataBaseConfig().getDatabaseName();
        if (databaseName == null) {
            databaseName = this.connDatabaseName;
        }

        String schema = generatorConfig.getDataBaseConfig().getSchema();

        schema = Objects.isNull(schema) ? connSchema : schema;

        if (generatorConfig.getTableConfig().getIncludeTables().isEmpty() || generatorConfig.getTableConfig().getIncludeTables().size() > 5) {
            return this.getTableInfoList(databaseName, schema, null, types);
        } else {
            List tables = new ArrayList<>();
            for (String table : generatorConfig.getTableConfig().getIncludeTables()) {
                List list = this.getTableInfoList(databaseName, schema, table, types);
                if (list.isEmpty()) {
                    list = this.getTableInfoList(databaseName, schema, table.toUpperCase(), types);
                }
                tables.addAll(list);
            }
            return tables;
        }
    }

    private List getColumnInfo(TableInfo tableInfo, String tableName) {
        Set primaryKeys = new HashSet<>();
        if (!tableInfo.isView()) {
            try (ResultSet primaryKeysResultSet = metaData.getPrimaryKeys(tableInfo.getCatalog(), tableInfo.getSchema(), tableName)) {
                while (primaryKeysResultSet.next()) {
                    String columnName = primaryKeysResultSet.getString("COLUMN_NAME");
                    if (columnName.toUpperCase().equals(columnName)) {
                        columnName = columnName.toLowerCase();
                    }
                    primaryKeys.add(columnName);
                }
            } catch (SQLException e) {
                throw new RuntimeException("读取表主键信息:" + tableInfo.getName() + "错误:", e);
            }
        }

        if (primaryKeys.isEmpty()) {
            log.warn("当前表:{}，存在无主键情况！", tableInfo.getName());
        }

        List columnsInfoList = new ArrayList<>();
        try (ResultSet resultSet = metaData.getColumns(tableInfo.getCatalog(), tableInfo.getSchema(), tableName, "%")) {
            while (resultSet.next()) {
                String columnName = resultSet.getString("COLUMN_NAME");
                if (columnName.toUpperCase().equals(columnName)) {
                    columnName = columnName.toLowerCase();
                }

                ColumnInfo columnInfo = new ColumnInfo();
                columnInfo.setTableInfo(tableInfo);
                columnInfo.setName(columnName);
                columnInfo.setPrimaryKey(primaryKeys.contains(columnInfo.getName()));

                columnInfo.setTypeName(resultSet.getString("TYPE_NAME"));
                columnInfo.setJdbcType(JdbcType.forCode(resultSet.getInt("DATA_TYPE")));
                columnInfo.setLength(resultSet.getInt("COLUMN_SIZE"));
                columnInfo.setScale(resultSet.getInt("DECIMAL_DIGITS"));
                columnInfo.setRemarks(resultSet.getString("REMARKS"));


                columnInfo.setDefaultValue(generatorConfig.getColumnConfig().getDefaultValueConvert().apply(resultSet.getString("COLUMN_DEF")));
                columnInfo.setNullable(resultSet.getInt("NULLABLE") == DatabaseMetaData.columnNullable);

                columnInfo.setVersion(columnName.equals(generatorConfig.getColumnConfig().getVersionColumn()));
                columnInfo.setTenantId(columnName.equals(generatorConfig.getColumnConfig().getTenantIdColumn()));
                columnInfo.setLogicDelete(columnName.equals(generatorConfig.getColumnConfig().getLogicDeleteColumn()));
                try {
                    columnInfo.setAutoIncrement("YES".equals(resultSet.getString("IS_AUTOINCREMENT")));
                } catch (SQLException e) {
                    log.error("获取自增信息异常：", e);
                }
                columnsInfoList.add(columnInfo);
            }
            return columnsInfoList;
        } catch (SQLException e) {
            throw new RuntimeException("读取表字段信息:" + tableInfo.getName() + "错误:", e);
        }
    }
}