1. 程式人生 > 實用技巧 >DBUnit備份資料庫,表,資料還原

DBUnit備份資料庫,表,資料還原

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" );


    }

}