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