tiesheng-springboot/springboot-database/src/main/java/com/tiesheng/database/service/DbMigrationInitializer.java

package com.tiesheng.database.service;


import cn.hutool.core.date.DateUtil;
import cn.hutool.core.io.IoUtil;
import cn.hutool.core.lang.func.VoidFunc1;
import cn.hutool.core.util.IdUtil;
import cn.hutool.core.util.StrUtil;
import cn.hutool.db.Db;
import cn.hutool.db.Entity;
import com.tiesheng.database.config.DbMigrationConfig;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.boot.web.servlet.ServletContextInitializer;
import org.springframework.core.io.Resource;
import org.springframework.core.io.support.PathMatchingResourcePatternResolver;
import org.springframework.stereotype.Component;

import javax.servlet.ServletContext;
import javax.sql.DataSource;
import java.io.InputStream;
import java.sql.SQLException;
import java.util.ArrayList;
import java.util.Arrays;
import java.util.List;
import java.util.concurrent.atomic.AtomicReference;
import java.util.zip.Checksum;

/**
 * 数据库表版本对比
 *
 * @author hao
 */
@Component
public class DbMigrationInitializer implements ServletContextInitializer {

    @Autowired
    DataSource dataSource;
    @Autowired
    DbMigrationConfig dbMigrationConfig;
    @Autowired
    List<DbDataInitializer> dbDataInitializerList;


    @Override
    public void onStartup(ServletContext servletContext) {
        try {

            // 初始化
            Db coreDb = Db.use(dataSource);
            dbMigrationConfig.checkDbExists(coreDb);

            startDeal(coreDb, dbMigrationConfig.getMigrations());
            startDeal(coreDb, dbMigrationConfig.getDatas());

            // 其他数据初始化
            dbDataInitializerList.forEach(DbDataInitializer::init);

        } catch (Exception ignore) {
        }
    }


    /**
     * 开始处理数据
     *
     * @throws Exception
     */
    private void startDeal(Db coreDb, List<String> locations) throws Exception {
        PathMatchingResourcePatternResolver patternResolver = new PathMatchingResourcePatternResolver();
        List<Resource> resourceList = new ArrayList<>();
        for (String location : locations) {
            Resource[] resources = patternResolver.getResources(location);
            resourceList.addAll(Arrays.asList(resources));
        }
        resourceList.sort((o1, o2) -> StrUtil.compare(o1.getFilename(), o2.getFilename(), true));
        for (Resource resource : resourceList) {
            migrationByResource(resource, coreDb);
        }
    }


    /**
     * 根据资源合并数据库结构
     *
     * @param resource
     * @param db
     * @throws Exception
     */
    private void migrationByResource(Resource resource, Db db) throws Exception {
        InputStream inputStream = resource.getInputStream();
        String readUtf8 = IoUtil.readUtf8(inputStream);
        String filename = resource.getFilename();
        Checksum checksum = IoUtil.checksum(resource.getInputStream(), null);
        Entity entity = fileChecksum(db, filename, checksum.getValue());
        if (entity == null) {
            return;
        }
        AtomicReference<Integer> success = new AtomicReference<>(0);
        AtomicReference<Integer> fail = new AtomicReference<>(0);
        db.tx((VoidFunc1<Db>) parameter -> {
            List<String> split = StrUtil.split(readUtf8, ";");
            for (String sql : split) {
                sql = sql.trim();
                if (StrUtil.isEmpty(sql)) {
                    continue;
                }
                if (StrUtil.containsAnyIgnoreCase(sql, StrUtil.splitToArray(dbMigrationConfig.getIgnoreSqls(), ","))) {
                    continue;
                }
                try {
                    parameter.execute(sql);
                    success.getAndSet(success.get() + 1);
                } catch (Exception ignore) {
                    fail.getAndSet(fail.get() + 1);
                }
            }
        });
        entity.set("checksum", checksum.getValue());
        entity.set("update_time", DateUtil.date());
        db.update(entity, Entity.create(dbMigrationConfig.getTable()).set("id", entity.get("id")));
    }


    /**
     * 检查文件是否执行
     *
     * @param db
     * @param fileName
     * @param checksum
     */
    private Entity fileChecksum(Db db, String fileName, long checksum) throws SQLException {
        Entity entity = db.get(dbMigrationConfig.getTable(), "file", fileName);
        if (entity == null) {
            entity = new Entity(dbMigrationConfig.getTable());
            entity.set("id", IdUtil.getSnowflakeNextId());
            entity.set("create_time", DateUtil.date());
            entity.set("update_time", DateUtil.date());
            entity.set("is_deleted", 0);
            entity.set("file", fileName);
            entity.set("checksum", 0);
            db.insert(entity);
        }

        if (entity.getLong("checksum") != checksum) {
            return entity;
        }
        return null;
    }


}