com.github.dactiv.common.unit.Fixtures.java Source code

Introduction

Here is the source code for com.github.dactiv.common.unit.Fixtures.java

Source

/**
 * Copyright (c) 2005-2012 springside.org.cn
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 */
package com.github.dactiv.common.unit;

import java.io.IOException;
import java.io.InputStream;
import java.sql.ResultSet;
import java.sql.SQLException;
import java.util.ArrayList;
import java.util.Arrays;
import java.util.HashMap;
import java.util.List;

import javax.sql.DataSource;

import org.dbunit.DatabaseUnitException;
import org.dbunit.database.IDatabaseConnection;
import org.dbunit.dataset.IDataSet;
import org.dbunit.dataset.xml.FlatXmlDataSetBuilder;
import org.dbunit.ext.h2.H2Connection;
import org.dbunit.operation.DatabaseOperation;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.core.io.DefaultResourceLoader;
import org.springframework.core.io.ResourceLoader;
import org.springframework.jdbc.core.namedparam.NamedParameterJdbcTemplate;

/**
 * DBUnit??H2?.
 * 
 * @author calvin
 */
public abstract class Fixtures {

    private static Logger logger = LoggerFactory.getLogger(Fixtures.class);
    private static ResourceLoader resourceLoader = new DefaultResourceLoader();

    /**
     * ?XML?H2?.
     *  
     * @param xmlFilePaths ?Spring Resource?.
     */
    public static void loadData(DataSource h2DataSource, String... xmlFilePaths) throws Exception {

        execute(DatabaseOperation.INSERT, h2DataSource, xmlFilePaths);
    }

    /**
     * XML???, ??XML?H2?.
     * 
     * reloadAllTable?, ????.
     * 
     * @param xmlFilePaths ?Spring Resource?.
     */
    public static void reloadData(DataSource h2DataSource, String... xmlFilePaths) throws Exception {
        execute(DatabaseOperation.CLEAN_INSERT, h2DataSource, xmlFilePaths);
    }

    /**
     * ??, ??XML?H2?.
     * 
     * @param xmlFilePaths ?Spring Resource?.
     */
    public static void reloadAllTable(DataSource h2DataSource, String... xmlFilePaths) throws Exception {
        deleteAllTable(h2DataSource);
        loadData(h2DataSource, xmlFilePaths);
    }

    /**
     * H2?XML??. 
     * 
     * @param xmlFilePaths ?Spring Resource?.
     */
    public static void deleteData(DataSource h2DataSource, String... xmlFilePaths) throws Exception {
        execute(DatabaseOperation.DELETE_ALL, h2DataSource, xmlFilePaths);
    }

    /**
     * XML?H2?Operation.
     * 
     * @param xmlFilePaths ?Spring Resource?.
     */
    private static void execute(DatabaseOperation operation, DataSource dataSource, String... xmlFilePaths)
            throws DatabaseUnitException, SQLException {
        //?HardCodeH2Connetion
        IDatabaseConnection connection = new H2Connection(dataSource.getConnection(), null);

        for (String xmlPath : xmlFilePaths) {
            try {
                InputStream input = resourceLoader.getResource(xmlPath).getInputStream();
                IDataSet dataSet = new FlatXmlDataSetBuilder().setColumnSensing(true).build(input);
                operation.execute(connection, dataSet);
            } catch (IOException e) {
                logger.warn(xmlPath + " file not found", e);
            } finally {
                connection.close();
            }
        }
    }

    /**
     * ,excludeTables.disable.
     * @throws SQLException 
     */
    public static void deleteAllTable(DataSource h2DataSource, String... excludeTables) throws SQLException {

        List<String> tableNames = new ArrayList<String>();

        try {
            ResultSet rs = h2DataSource.getConnection().getMetaData().getTables(null, null, null,
                    new String[] { "TABLE" });

            while (rs.next()) {
                String tableName = rs.getString("TABLE_NAME");
                if (Arrays.binarySearch(excludeTables, tableName) < 0) {
                    tableNames.add(tableName);
                }
            }

            deleteTable(h2DataSource, tableNames.toArray(new String[tableNames.size()]));

        } catch (SQLException e) {
            throw e;
        }

    }

    /**
     * , disable.
     */
    public static void deleteTable(DataSource h2DataSource, String... tableNames) {
        NamedParameterJdbcTemplate template = new NamedParameterJdbcTemplate(h2DataSource);

        template.update("SET REFERENTIAL_INTEGRITY FALSE", new HashMap<String, Object>());

        for (String tableName : tableNames) {
            template.update("DELETE FROM " + tableName, new HashMap<String, Object>());
        }

        template.update("SET REFERENTIAL_INTEGRITY TRUE", new HashMap<String, Object>());
    }
}