java实现数据库自检缺失表及字段(识别@TableName)

import cn.hutool.core.text.CharSequenceUtil;
import cn.hutool.core.util.ObjectUtil;
import com.baomidou.mybatisplus.annotation.TableField;
import com.baomidou.mybatisplus.annotation.TableName;
import com.fasterxml.jackson.databind.ObjectMapper;
import com.zaxxer.hikari.HikariDataSource;
import lombok.extern.slf4j.Slf4j;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.context.annotation.ClassPathScanningCandidateComponentProvider;
import org.springframework.core.env.Environment;
import org.springframework.core.type.filter.AnnotationTypeFilter;
import org.springframework.scheduling.concurrent.ThreadPoolTaskExecutor;
import org.springframework.stereotype.Component;

import javax.annotation.Resource;
import javax.sql.DataSource;
import java.io.File;
import java.io.FileWriter;
import java.io.IOException;
import java.lang.reflect.Field;
import java.sql.Connection;
import java.sql.DatabaseMetaData;
import java.sql.ResultSet;
import java.sql.SQLException;
import java.util.*;
import java.util.concurrent.CompletableFuture;
import java.util.concurrent.ConcurrentHashMap;

@Slf4j
@Component
public class DbScriptChecker {

    @Autowired
    private Environment env;

    @Resource
    private ThreadPoolTaskExecutor executor;

    private final Map<String, List<Class<?>>> scannedClassesCache = new ConcurrentHashMap<>();
    private final Map<String, Set<String>> tableColumnsCache = new ConcurrentHashMap<>();

    public void checkDatabaseSchema() {
        log.info("数据库自检::开始进行数据库表匹配自检功能");
        DataSource dataSource = this.createDataSource();
        List<String> scanPackages = Arrays.asList(env.getProperty("db.script.checker.scan.packages").split(","));
        List<Class<?>> tableClasses = new ArrayList<>();
        for (String packageName : scanPackages) {
            tableClasses.addAll(this.scanForTableClasses(packageName));
        }
        if(ObjectUtil.isEmpty(tableClasses)){
            log.error("数据库自检::未找到任何表类, 请检查配置中的包路径是否正确");
            return;
        }

        Map<String, List<String>> missingTables = new HashMap<>();
        Map<String, List<String>> missingFields = new HashMap<>();

        try (Connection connection = dataSource.getConnection()) {
            DatabaseMetaData metaData = connection.getMetaData();

            // 使用 CompletableFuture 并行处理每个表的检查
            List<CompletableFuture<?>> futures = new ArrayList<>();
            for (Class<?> tableClass : tableClasses) {
                CompletableFuture<Void> future = CompletableFuture.runAsync(() -> {
                    TableName tableNameAnnotation = tableClass.getAnnotation(TableName.class);
                    if (tableNameAnnotation == null) {
                        log.warn("数据库自检::类: {} 没有 @TableName 注解", tableClass.getName());
                        return;
                    }

                    String tableName = tableNameAnnotation.value();
                    log.info("数据库自检::正在检查{}表", tableName);

                    // 检查表是否存在
                    boolean tableExists = false;
                    try {
                        tableExists = this.tableExists(metaData, tableName);
                    } catch (SQLException e) {
                        log.error("数据库自检::表 {} 异常", tableName);
                        return;
                    }
                    if (!tableExists) {
                        log.error("数据库自检::表 {} 在数据库中不存在", tableName);
                        missingTables.computeIfAbsent(tableName, k -> new ArrayList<>()).add(tableClass.getName());
                        return;
                    }

                    // 获取表的字段
                    Set<String> dbColumns = null;
                    try {
                        dbColumns = this.getTableColumns(metaData, tableName);
                    } catch (SQLException e) {
                        log.error("数据库自检::表{}.字段{} 异常", tableName, null);
                        return;
                    }

                    // 检查字段是否存在
                    List<Field> fields = Arrays.asList(tableClass.getDeclaredFields());
                    for (Field field : fields) {
                        if ("serialVersionUID".equals(field.getName())) {
                            continue;
                        }

                        TableField tableFieldAnnotation = field.getAnnotation(TableField.class);
                        if (tableFieldAnnotation != null && !tableFieldAnnotation.exist()) {
                            continue;
                        }

                        String columnName = CharSequenceUtil.toUnderlineCase(field.getName());
                        if (!dbColumns.contains(columnName)) {
                            log.error("数据库自检::字段 {} 不在表中 {}", columnName, tableName);
                            missingFields.computeIfAbsent(tableName, k -> new ArrayList<>()).add(columnName);
                        }
                    }
                }, executor);
                futures.add(future);
            }

            // 等待所有任务完成
            CompletableFuture.allOf(futures.toArray(new CompletableFuture[0])).join();
        } catch (SQLException e) {
            log.error("数据库自检::数据库表自检失败: {}", e.getMessage());
            return;
        }

        this.generateJsonReport(missingTables, missingFields);

        log.info("数据库自检::自检完毕.");
        executor.shutdown();
    }

    private DataSource createDataSource() {
        log.info("数据库自检::创建数据库连接...");
        HikariDataSource dataSource = new HikariDataSource();
        dataSource.setDriverClassName(env.getProperty("db.script.checker.datasource.driver-class-name"));
        dataSource.setJdbcUrl(env.getProperty("db.script.checker.datasource.url"));
        dataSource.setUsername(env.getProperty("db.script.checker.datasource.username"));
        dataSource.setPassword(env.getProperty("db.script.checker.datasource.password"));
        dataSource.setMaximumPoolSize(10);
        dataSource.setMinimumIdle(2);
        return dataSource;
    }

    private boolean tableExists(DatabaseMetaData metaData, String tableName) throws SQLException {
        try (ResultSet tables = metaData.getTables(null, null, tableName, new String[]{"TABLE"})) {
            return tables.next();
        }
    }

    private Set<String> getTableColumns(DatabaseMetaData metaData, String tableName) throws SQLException {
        if (tableColumnsCache.containsKey(tableName)) {
            return tableColumnsCache.get(tableName);
        }

        Set<String> columns = new HashSet<>();
        try (ResultSet rs = metaData.getColumns(null, null, tableName, null)) {
            while (rs.next()) {
                columns.add(rs.getString("COLUMN_NAME"));
            }
        }
        tableColumnsCache.put(tableName, columns);
        return columns;
    }

    private List<Class<?>> scanForTableClasses(String packageName) {
        if (scannedClassesCache.containsKey(packageName)) {
            return scannedClassesCache.get(packageName);
        }

        log.info("数据库自检::扫描的包: {}", packageName);
        List<Class<?>> tableClasses = new ArrayList<>();
        ClassPathScanningCandidateComponentProvider scanner = new ClassPathScanningCandidateComponentProvider(false);
        scanner.addIncludeFilter(new AnnotationTypeFilter(TableName.class));
        for (org.springframework.beans.factory.config.BeanDefinition beanDefinition : scanner.findCandidateComponents(packageName)) {
            try {
                Class<?> clazz = Class.forName(beanDefinition.getBeanClassName());
                tableClasses.add(clazz);
                log.info("数据库自检::查找到 @TableName 注解: {}", clazz.getName());
            } catch (ClassNotFoundException e) {
                log.error("数据库自检::没有找到: {}", beanDefinition.getBeanClassName(), e);
            }
        }
        scannedClassesCache.put(packageName, tableClasses);
        return tableClasses;
    }

    private void generateJsonReport(Map<String, List<String>> missingTables, Map<String, List<String>> missingFields) {
        log.info("数据库自检::开始生成 JSON 报表...");
        Map<String, Object> report = new HashMap<>();
        report.put("缺失表", missingTables);
        report.put("缺失字段", missingFields);

        ObjectMapper objectMapper = new ObjectMapper();
        try {
            String jsonReport = objectMapper.writerWithDefaultPrettyPrinter().writeValueAsString(report);
            File jsonFile = new File("db_script_check_report.json");
            try (FileWriter fileWriter = new FileWriter(jsonFile)) {
                fileWriter.write(jsonReport);
            }
            log.info("数据库自检::JSON 报表生成完毕: {}", jsonFile.getAbsolutePath());
        } catch (IOException e) {
            log.error("数据库自检::JSON 报表生成错误: {}", e.getMessage());
        }
    }
}