package com.util;

import com.config.GetTestProperties;
import org.apache.log4j.Logger;
import org.apache.log4j.PropertyConfigurator;
import org.dbunit.DBTestCase;
import org.dbunit.PropertiesBasedJdbcDatabaseTester;
import org.dbunit.database.IDatabaseConnection;
import org.dbunit.database.QueryDataSet;
import org.dbunit.dataset.IDataSet;
import org.dbunit.dataset.xml.FlatXmlDataSet;
import org.dbunit.operation.DatabaseOperation;

import java.io.File;
import java.io.FileInputStream;
import java.io.FileOutputStream;
import java.sql.Connection;
import java.sql.DriverManager;
import java.sql.SQLException;
import java.util.List;

// 数据库备份、恢复服务
public class DBUnit extends DBTestCase {
    private final Logger log = Logger.getLogger( DBUnit.class.getClass().getName() );
    private String dir_name;
    private DBUnit() {
        dir_name = GetTestProperties.getDbBackup();
        String dbType = GetTestProperties.getDatabasesType();
        String dbUsername = GetTestProperties.getDatabasesUsername();
        String dbPassword = GetTestProperties.getDatabasesPassword();
        String dbUrl = GetTestProperties.getDatabasesUrl();

        try {
            PropertyConfigurator.configure( CommonMethord.getRealath()
                    + "log4j.properties" );
        } catch (Exception e) {
            e.printStackTrace();
        }
        CommonMethord.createDir( dir_name );
        if (dbType.equals( "mysql" )) {
            System.setProperty(
                    PropertiesBasedJdbcDatabaseTester.DBUNIT_DRIVER_CLASS,
                    "com.mysql.jdbc.Driverc" );
        } else if (dbType.equals( "oracle" )) {
            System.setProperty(
                    PropertiesBasedJdbcDatabaseTester.DBUNIT_DRIVER_CLASS,
                    "oracle.jdbc.driver.OracleDriver" );
        } else {
            log.error( "未定义的数据库类型 !" );
        }

        System.setProperty(
                PropertiesBasedJdbcDatabaseTester.DBUNIT_CONNECTION_URL, dbUrl );
        System.setProperty( PropertiesBasedJdbcDatabaseTester.DBUNIT_USERNAME,
                dbUsername );
        System.setProperty( PropertiesBasedJdbcDatabaseTester.DBUNIT_PASSWORD,
                dbPassword );

    }


    /**
     * 给定数据集
     *
     * @return
     * @throws Exception
     */
    @SuppressWarnings("deprecation")
    @Override
    protected IDataSet getDataSet() throws Exception {
        log.info( "init..." );
        return new FlatXmlDataSet( new FileInputStream( "" ) );
    }

    /**
     * getSetUpOperation
     *
     * @return
     * @throws Exception
     */
    @Override
    protected DatabaseOperation getSetUpOperation() throws Exception {
        return DatabaseOperation.REFRESH;
    }

    /**
     * getTearDownOperation
     *
     * @return
     * @throws Exception
     */
    @Override
    protected DatabaseOperation getTearDownOperation() throws Exception {
        return DatabaseOperation.NONE;
    }

    /**
     * 给定数据集
     *
     * @param fileName
     * @return
     * @throws Exception
     */
    @SuppressWarnings("deprecation")
    protected IDataSet getDataSet(String fileName) throws Exception {
        log.info( "init..." );
        return new FlatXmlDataSet( new FileInputStream( fileName ) );
    }

    /**
     * 单表备份,传入表名和备份文件名
     *
     * @param tbname
     * @param xmlFileName
     * @throws Exception
     */
    public void backupTable(String tbname, String xmlFileName) throws Exception {
        IDatabaseConnection connection = getConnection();
        try {

            QueryDataSet dataSet = new QueryDataSet( connection );
            // 将表里的数据导出到 xml文件里
            dataSet.addTable( tbname );
            // 将表里符合条件的数据导出到xml文件里
            // dataSet.addTable("users", "select * from users where id < 4");
            // 导出到dbunit.xml文件里
            File file = new File( dir_name + File.separator + xmlFileName );
            FlatXmlDataSet.write( dataSet, new FileOutputStream( file ) );
            log.info( "单表备份完成！" );
        } catch (Exception e) {
            log.error( "无法连接到数据库服务器" );
            e.printStackTrace();
        } finally {
            try {
                if (connection != null)
                    connection.close();
            } catch (SQLException e) {
                e.printStackTrace();
            }
        }
    }

    /**
     * 备份多个数据表 ,英文逗号分隔多个表名
     *
     * @param tabname
     * @param xmlFileName
     * @throws Exception
     */
    public void backupTables(String tabname, String xmlFileName) throws Exception {
        IDatabaseConnection connection = getConnection();
        try {
            String tbname;
            List<String> tbs = CommonMethord.getList( tabname );
            QueryDataSet dataSet = new QueryDataSet( connection );
            // 添加多个table
            for (int i = 0; i < tbs.size(); i++) {
                tbname = tbs.get( i );
                dataSet.addTable( tbname );
            }
            // 导出到dbunit.xml文件里
            File f_file = new File( dir_name + File.separator + xmlFileName );
            FlatXmlDataSet.write( dataSet, new FileOutputStream( f_file ) );
            log.info( "多表备份完成！" );
        } catch (Exception e) {
            log.error( "无法连接到数据库服务器" );
            e.printStackTrace();
        } finally {
            try {
                if (connection != null)
                    connection.close();
            } catch (SQLException e) {
                e.printStackTrace();
            }
        }
    }

    /**
     * 备份全部数据表
     *
     * @param xmlFileName
     * @throws Exception
     */
    public void backupAllTables(String xmlFileName) throws Exception {
        IDatabaseConnection connection = getConnection();

        try {
            // 如果想把某个数据库里的所有表里的数据全部导出到某个xml里,又不想通过addTable一个个来添加的话。
            // 则必须通过IDatabaseConnection的createDataSet()来创建IDataSet
            IDataSet dataSet = connection.createDataSet();
            // 导出到dbunit.xml文件里
            File f_file = new File( dir_name + File.separator + xmlFileName );
            FlatXmlDataSet.write( dataSet, new FileOutputStream( f_file ) );
            // 也可以用FlatDtdDataSet导出一个对应的dtd文件
            // FlatDtdDataSet.write(dataSet, new FileOutputStream(
            // "dbunit_alltb.dtd"));

        } catch (Exception e) {
            log.error( "无法连接到数据库服务器" );
            e.printStackTrace();
        } finally {
            try {
                if (connection != null)
                    connection.close();
            } catch (SQLException e) {
                e.printStackTrace();
            }
        }
    }

    /**
     * restoreDb 恢复DB, 在Maven恢复较慢
     *
     * @param xmlFileName
     * @throws Exception
     */
    public void restoreDb(String xmlFileName) throws Exception {
        IDatabaseConnection connection = getConnection();
        log.info( "连接成功！" );
        try {
            //IDataSet xmlDataSet = new FlatXmlDataSet( new FileInputStream( xmlFileName ) );
            File file = new File( dir_name + File.separator + xmlFileName );
            IDataSet xmlDataSet = getDataSet( file.getAbsolutePath() );
            DatabaseOperation.CLEAN_INSERT.execute( connection, xmlDataSet );
            log.info( "数据恢复完成！" );
        } catch (Exception e) {
            log.error( "无法连接到数据库服务器" );
            e.printStackTrace();
        } finally {
            try {
                if (connection != null)
                    connection.close();
            } catch (SQLException e) {
                e.printStackTrace();
            }
        }
    }




        // 测试
    public static void main(String[] args) throws Exception {

        DBUnit db = new DBUnit();
        // 单表
        db.backupTable( "LOG_DATA_01", "LOG_DATA_01" );
        // 多表
//        db.backupTables( "MENU,P_EAST_LOG,RESMODIFYLOG,ROLE_RIGHT,RULE_TABLE","all" );
        // 全表
//        db.backupAllTables( "all_table" );

        // 恢复数据
//        db.restoreDb( "UMG_USER" );


    }

}