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 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 locations) throws Exception { PathMatchingResourcePatternResolver patternResolver = new PathMatchingResourcePatternResolver(); List 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 success = new AtomicReference<>(0); AtomicReference fail = new AtomicReference<>(0); db.tx((VoidFunc1) parameter -> { List 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; } }