001package cn.mybatis.mp.generator.database.meta;
002
003import cn.mybatis.mp.generator.config.GeneratorConfig;
004import lombok.extern.slf4j.Slf4j;
005import org.apache.ibatis.type.JdbcType;
006
007import java.sql.Connection;
008import java.sql.DatabaseMetaData;
009import java.sql.ResultSet;
010import java.sql.SQLException;
011import java.util.ArrayList;
012import java.util.HashSet;
013import java.util.List;
014import java.util.Set;
015import java.util.stream.Collectors;
016
017@Slf4j
018public class TableMetaDataQuery {
019
020    protected final DatabaseMetaData metaData;
021
022    protected final GeneratorConfig generatorConfig;
023
024    protected final String connDatabaseName;
025
026    public TableMetaDataQuery(GeneratorConfig generatorConfig, Connection connection) throws SQLException {
027        this.metaData = connection.getMetaData();
028        this.generatorConfig = generatorConfig;
029        this.connDatabaseName = connection.getCatalog();
030    }
031
032    public List<TableInfo> getTableInfoList(boolean includeTable, boolean includeView) {
033        Set<String> types = new HashSet<>();
034        if (includeTable) {
035            types.add("TABLE");
036        }
037        if (includeView) {
038            types.add("VIEW");
039        }
040
041        List<TableInfo> tables = new ArrayList<>();
042
043        String databaseName = generatorConfig.getDataBaseConfig().getDatabaseName();
044        if (databaseName == null) {
045            databaseName = this.connDatabaseName;
046        }
047
048        try (ResultSet resultSet = metaData.getTables(databaseName, generatorConfig.getDataBaseConfig().getSchema(), null, types.toArray(new String[2]))) {
049            TableInfo tableInfo;
050            while (resultSet.next()) {
051                String TABLE_NAME = resultSet.getString("TABLE_NAME");
052                String tableName = TABLE_NAME;
053                if (tableName.toUpperCase().equals(tableName)) {
054                    tableName = tableName.toLowerCase();
055                }
056                tableInfo = new TableInfo();
057                tableInfo.setName(tableName);
058                tableInfo.setRemarks(resultSet.getString("REMARKS"));
059                tableInfo.setTableType(resultSet.getString("TABLE_TYPE"));
060                tableInfo.setSchema(resultSet.getString("TABLE_SCHEM"));
061                tableInfo.setCatalog(resultSet.getString("TABLE_CAT"));
062
063                tableInfo.setColumnInfoList(getColumnInfo(tableInfo, TABLE_NAME));
064                List<ColumnInfo> idColumnInfoList = tableInfo.getColumnInfoList().stream().filter(item -> item.isPrimaryKey()).collect(Collectors.toList());
065                if (!idColumnInfoList.isEmpty()) {
066                    tableInfo.setIdColumnInfo(idColumnInfoList.get(0));
067                }
068                tables.add(tableInfo);
069            }
070        } catch (SQLException e) {
071            throw new RuntimeException("读取数据库表信息出现错误", e);
072        }
073        return tables;
074    }
075
076    private List<ColumnInfo> getColumnInfo(TableInfo tableInfo, String tableName) {
077        Set<String> primaryKeys = new HashSet<>();
078        if (!tableInfo.isView()) {
079            try (ResultSet primaryKeysResultSet = metaData.getPrimaryKeys(tableInfo.getCatalog(), tableInfo.getSchema(), tableName)) {
080                while (primaryKeysResultSet.next()) {
081                    String columnName = primaryKeysResultSet.getString("COLUMN_NAME");
082                    if (columnName.toUpperCase().equals(columnName)) {
083                        columnName = columnName.toLowerCase();
084                    }
085                    primaryKeys.add(columnName);
086                }
087            } catch (SQLException e) {
088                throw new RuntimeException("读取表主键信息:" + tableInfo.getName() + "错误:", e);
089            }
090        }
091
092        if (primaryKeys.size() > 1) {
093            log.warn("当前表:{}，存在多主键情况！", tableInfo.getName());
094        }
095
096        List<ColumnInfo> columnsInfoList = new ArrayList<>();
097        try (ResultSet resultSet = metaData.getColumns(tableInfo.getCatalog(), tableInfo.getSchema(), tableName, "%")) {
098            while (resultSet.next()) {
099                String columnName = resultSet.getString("COLUMN_NAME");
100                if (columnName.toUpperCase().equals(columnName)) {
101                    columnName = columnName.toLowerCase();
102                }
103
104                ColumnInfo columnInfo = new ColumnInfo();
105                columnInfo.setTableInfo(tableInfo);
106                columnInfo.setName(columnName);
107                columnInfo.setPrimaryKey(primaryKeys.contains(columnInfo.getName()));
108
109                columnInfo.setTypeName(resultSet.getString("TYPE_NAME"));
110                columnInfo.setJdbcType(JdbcType.forCode(resultSet.getInt("DATA_TYPE")));
111                columnInfo.setLength(resultSet.getInt("COLUMN_SIZE"));
112                columnInfo.setScale(resultSet.getInt("DECIMAL_DIGITS"));
113                columnInfo.setRemarks(resultSet.getString("REMARKS"));
114                columnInfo.setDefaultValue(resultSet.getString("COLUMN_DEF"));
115                columnInfo.setNullable(resultSet.getInt("NULLABLE") == DatabaseMetaData.columnNullable);
116
117                columnInfo.setVersion(columnName.equals(generatorConfig.getColumnConfig().getVersionColumn()));
118                columnInfo.setTenantId(columnName.equals(generatorConfig.getColumnConfig().getTenantIdColumn()));
119                try {
120                    columnInfo.setAutoIncrement("YES".equals(resultSet.getString("IS_AUTOINCREMENT")));
121                } catch (SQLException e) {
122                    log.error("获取自增信息异常：", e);
123                }
124                columnsInfoList.add(columnInfo);
125            }
126            return columnsInfoList;
127        } catch (SQLException e) {
128            throw new RuntimeException("读取表字段信息:" + tableInfo.getName() + "错误:", e);
129        }
130    }
131}