org.test4j.module.database.proxy.DataSourceScriptHelper Maven / Gradle / Ivy

Show more of this group Show more artifacts with this name
Show all versions of fluent-dbtest Show documentation
fluent database test
There is a newer version: 1.1.2
package org.test4j.module.database.proxy;

import cn.org.atool.fluent.mybatis.metadata.DbType;
import org.springframework.core.io.ClassPathResource;
import org.springframework.core.io.InputStreamResource;
import org.springframework.core.io.Resource;
import org.springframework.core.io.support.EncodedResource;
import org.springframework.jdbc.datasource.DataSourceUtils;
import org.test4j.hamcrest.IWant;
import org.test4j.hamcrest.base.Assert;
import org.test4j.module.database.IDataSourceScript;
import org.test4j.module.database.config.DbConfig;
import org.test4j.tools.commons.StringHelper;

import javax.sql.DataSource;
import java.io.ByteArrayInputStream;
import java.sql.Connection;
import java.sql.SQLException;
import java.util.ArrayList;
import java.util.HashMap;
import java.util.List;
import java.util.Map;

import static org.springframework.jdbc.datasource.init.ScriptUtils.executeSqlScript;

/**
 * @author wudarui
 */
public class DataSourceScriptHelper {
    /**
     * key: dataSource name
     */
    static Map DATASOURCE_SCRIPT_HAS_INIT = new HashMap<>();

    public static void runInitScripts(DataSource dataSource, String beanName) {
        runInitScripts(dataSource, beanName, null);
    }

    public static void runInitScripts(DataSource dataSource, String beanName, DbType dbType) {
        if (dataSource == null || DATASOURCE_SCRIPT_HAS_INIT.containsKey(beanName)) {
            return;
        }
        IDataSourceScript factory = DbConfig.instance().dbScriptFactory(beanName);
        boolean hasRun;
        if (factory == null) {
            hasRun = runFromClasspathResource(dataSource, beanName);
        } else {
            DbType type = dbType == null ? DbConfig.instance().dbType(beanName) : dbType;
            hasRun = runFromScriptFactory(type, dataSource, factory);
        }
        if (hasRun) {
            commitScript(dataSource);
        }
        DATASOURCE_SCRIPT_HAS_INIT.put(beanName, true);
    }

    private static void commitScript(DataSource dataSource) {
        try {
            if (!dataSource.getConnection().getAutoCommit()) {
                dataSource.getConnection().commit();
            }
        } catch (SQLException e) {
            e.printStackTrace();
            throw new RuntimeException("commit datasource script error:" + e.getMessage(), e);
        }
    }

    private static boolean runFromScriptFactory(DbType type, DataSource dataSource, IDataSourceScript factory) {
        String script = factory.script(type);
        try {
            IWant.want.string(script).notNull("script can't be null.");
            Connection connection = DataSourceUtils.getConnection(dataSource);
            try {
                Resource resource = new InputStreamResource(new ByteArrayInputStream(script.getBytes()));
                executeSqlScript(connection, new EncodedResource(resource, "utf-8"), false, true, "--", ";", "/*", "*/");
            } finally {
                DataSourceUtils.releaseConnection(connection, dataSource);
            }
            return true;
        } catch (Exception e) {
            throw new RuntimeException("Failed to execute database script:\n\t" + script + "\n\n", e);
        }
    }

    private static boolean runFromClasspathResource(DataSource dataSource, String beanName) {
        List resources = getResources(beanName);
        if (resources.isEmpty()) {
            return false;
        }
        Assert.notNull(dataSource, "DataSource must be provided.");
        try {
            Connection connection = DataSourceUtils.getConnection(dataSource);
            IWant.want.object(connection).notNull("connection must be null.");
            try {
                for (Resource resource : resources) {
                    executeSqlScript(connection, new EncodedResource(resource, "utf-8"), false, true, "--", ";", "/*", "*/");
                }
            } finally {
                DataSourceUtils.releaseConnection(connection, dataSource);
            }
            return true;
        } catch (Exception e) {
            throw new RuntimeException("Failed to execute database script", e);
        }
    }

    private static List getResources(String beanName) {
        String[] locations = DbConfig.instance().dbScript(beanName);
        List resources = new ArrayList<>();
        for (String location : locations) {
            if (StringHelper.isBlank(location)) {
                continue;
            }
            Resource resource = new ClassPathResource(location);
            if (resource.exists()) {
                resources.add(resource);
            } else {
                throw new RuntimeException("The specified resource[" + beanName + "] does not exist.");
            }
        }
        return resources;
    }
}