DatabaseClone.java

/*******************************************************************************
 * Copyright (c) 2010-2011 VIVO Harvester Team. For full list of contributors, please see the AUTHORS file provided.
 * All rights reserved.
 * This program and the accompanying materials are made available under the terms of the new BSD license which accompanies this distribution, and is available at http://www.opensource.org/licenses/bsd-license.html
 ******************************************************************************/
package org.vivoweb.harvester.util;

import java.io.IOException;
import java.io.InputStream;
import java.io.OutputStream;
import java.sql.Connection;
import java.sql.DriverManager;
import java.sql.ResultSet;
import java.sql.ResultSetMetaData;
import java.sql.SQLException;
import java.sql.Types;
import java.util.HashMap;
import java.util.HashSet;
import java.util.Map;
import java.util.Set;
import org.apache.commons.lang.StringUtils;
import org.dbunit.DatabaseUnitException;
import org.dbunit.database.DatabaseConfig;
import org.dbunit.database.DatabaseConnection;
import org.dbunit.database.IDatabaseConnection;
import org.dbunit.database.QueryDataSet;
import org.dbunit.dataset.DataSetException;
import org.dbunit.dataset.IDataSet;
import org.dbunit.dataset.xml.FlatDtdDataSet;
import org.dbunit.operation.DatabaseOperation;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.vivoweb.harvester.util.args.ArgDef;
import org.vivoweb.harvester.util.args.ArgList;
import org.vivoweb.harvester.util.args.ArgParser;
import org.vivoweb.harvester.util.args.UsageException;

/**
 * Clone a database from one jdbc connection to another
 * @author Christopher Haines (hainesc@ctrip.ufl.edu)
 */
public class DatabaseClone {
	/**
	 * SLF4J Logger
	 */
	private static Logger log = LoggerFactory.getLogger(DatabaseClone.class);
	/**
	 * database to export from
	 */
	private IDatabaseConnection db1;
	/**
	 * database to import to
	 */
	private IDatabaseConnection db2;
	/**
	 * list of tables to export (null exports all)
	 */
	private String[] tables;
	/**
	 * map of DBUnit features/properties to objects (Boolean,String,String[])
	 */
	private Map<String, ? extends Object> dbUnitFeatures;
	/**
	 * use this database state file as the input database
	 */
	private InputStream inFile;
	/**
	 * output the state of the input database in this file
	 */
	private OutputStream outFile;
	
	/**
	 * Constructor
	 * @param inputConn database to input from
	 * @param inputFile use this database state file as the input database
	 * @param outputConn database to output to
	 * @param outputFile output the state of the input database in this file
	 * @param tableNames list of tables to export (null exports all)
	 * @param dbUnitFeatures map of DBUnit features/properties to objects (Boolean,String,String[])
	 * @throws IOException error resolving file or connecting to database
	 */
	public DatabaseClone(Connection inputConn, String inputFile, Connection outputConn, String outputFile, String[] tableNames, Map<String, ? extends Object> dbUnitFeatures) throws IOException {
		if(inputConn != null) {
			try {
				this.db1 = new DatabaseConnection(inputConn);
			} catch(DatabaseUnitException e) {
				throw new IOException(e);
			}
		} else {
			this.db1 = null;
		}
		this.inFile = FileAide.getInputStream(inputFile);
		if(outputConn != null) {
			try {
				this.db2 = new DatabaseConnection(outputConn);
			} catch(DatabaseUnitException e) {
				throw new IOException(e);
			}
		} else {
			this.db2 = null;
		}
		this.outFile = FileAide.getOutputStream(outputFile, true);
		this.tables = tableNames;
		this.dbUnitFeatures = dbUnitFeatures;
		// Add value check info
		if((this.inFile == null) ^ (this.db1 != null)) {
			throw new IllegalArgumentException("Must provide one of input database or input file");
		}
		if((this.outFile == null && this.db2 == null)) {
			throw new IllegalArgumentException("Must provide an output database and/or output file");
		}
		if(this.inFile != null && this.outFile != null) {
			if(this.db2 == null) {
				throw new IllegalArgumentException("This tool should not be used to copy a database state file, please provide an output database");
			}
			log.warn("This tool should not be used to copy a database state file, ignoring outputFile");
			this.outFile = null;
		}
	}
	
	/**
	 * Constructor
	 * @param args commandline arguments
	 * @throws IOException error resolving file or connecting to database
	 * @throws UsageException user requested usage message
	 */
	public DatabaseClone(String[] args) throws IOException, UsageException {
		this(getParser().parse(args));
	}
	
	/**
	 * Constructor
	 * @param argList option set of parsed args
	 * @throws IOException error resolving file or connecting to database
	 */
	private DatabaseClone(ArgList argList) throws IOException {
		this(
			initDBConn("input", argList.get("inputDriver"), argList.get("inputConnection"), argList.get("inputUsername"), argList.get("inputPassword")),
			argList.get("inputFile"),
			initDBConn("output", argList.get("outputDriver"), argList.get("outputConnection"), argList.get("outputUsername"), argList.get("outputPassword")),
			argList.get("outputFile"),
			argList.getAll("tableName").toArray(new String[]{}),
			getDbUnitFeatures(argList.getValueMap("DBUnitFeature"), argList.getAll("validTableType").toArray(new String[]{}))
		);
	}
	
	/**
	 * Get the dbUnitFeatures in correct form
	 * @param dbUnitFeatures the existing features
	 * @param validTableTypes valid table types
	 * @return the correct format
	 */
	private static Map<String, ? extends Object> getDbUnitFeatures(Map<String, String> dbUnitFeatures, String... validTableTypes) {
		Map<String, Object> trueFeatures = new HashMap<String, Object>();
		for(String feature : dbUnitFeatures.keySet()) {
			String value = dbUnitFeatures.get(feature);
			if(value == null) {
				trueFeatures.put(feature, null);
				continue;
			}
			if(value.equalsIgnoreCase("true")) {
				trueFeatures.put(feature, Boolean.TRUE);
				continue;
			}
			if(value.equalsIgnoreCase("false")) {
				trueFeatures.put(feature, Boolean.FALSE);
				continue;
			}
			String[] tempArray;
			if((tempArray = value.split(",")).length > 1) {
				trueFeatures.put(feature, tempArray);
				continue;
			}
			Integer tempInt = null;
			try {
				tempInt = Integer.valueOf(value);
			} catch(NumberFormatException e) {
				//ignore
			}
			if(tempInt != null) {
				trueFeatures.put(feature, tempInt);
				continue;
			}
			Float tempFloat = null;
			try {
				tempFloat = Float.valueOf(value);
			} catch(NumberFormatException e) {
				//ignore
			}
			if(tempFloat != null) {
				trueFeatures.put(feature, tempFloat);
				continue;
			}
			Class<?> tempClass = null;
			try {
				tempClass = Class.forName(value);
			} catch(ClassNotFoundException e) {
				//ignore
			}
			if(tempClass != null) {
				trueFeatures.put(feature, tempClass);
				continue;
			}
			trueFeatures.put(feature, value);
		}
		if(validTableTypes != null && validTableTypes.length > 0) {
			String tableTypeList = StringUtils.join(validTableTypes, ",");
			log.warn("validTableTypes Parameter is now deprecated, to reproduce this effect add the following parameter:\n" +
				"<Param name=\"DBUnitFeature\">http://www.dbunit.org/properties/tableType="+tableTypeList+"</Param>\n" +
				"  or\n" +
				"--DBUnitFeature http://www.dbunit.org/properties/tableType="+tableTypeList
			);
			trueFeatures.put("http://www.dbunit.org/properties/tableType", tableTypeList);
		}
		return trueFeatures;
	}
	
	/**
	 * Initialize a Database Connection
	 * @param dbname the name for this database
	 * @param driver the jdbc driver
	 * @param connLine the jdbc connection line
	 * @param user the username
	 * @param pass the password
	 * @return the database connection
	 * @throws IOException error connecting to database
	 */
	private static Connection initDBConn(String dbname, String driver, String connLine, String user, String pass) throws IOException {
		if(driver == null) {
			log.debug("No "+dbname+"Driver provided, not using "+dbname+" database");
			return null;
		}
		if(connLine == null) {
			log.debug("No "+dbname+"Connection provided, not using "+dbname+" database");
			return null;
		}
		if(user == null) {
			log.debug("No "+dbname+"Username provided, not using "+dbname+" database");
			return null;
		}
		if(pass == null) {
			log.debug("No "+dbname+"Password provided, not using "+dbname+" database");
			return null;
		}
		try {
			Class.forName(driver);
		} catch(ClassNotFoundException e) {
			throw new IllegalArgumentException(e);
		}
		try {
			return DriverManager.getConnection(connLine, user, pass);
		} catch(SQLException e) {
			throw new IOException(e);
		}
	}
	
	/**
	 * Clone a database
	 * @throws SQLException error connecting
	 * @throws DatabaseUnitException error
	 * @throws IOException error resolving connections
	 */
	public void execute() throws SQLException, DatabaseUnitException, IOException {
		String[] tableTypes = new String[]{"TABLE"};
		if(this.dbUnitFeatures != null) {
			for(String feature : this.dbUnitFeatures.keySet()) {
				if(feature.trim().equals("http://www.dbunit.org/properties/tableType")) {
					tableTypes = this.dbUnitFeatures.get(feature).toString().split(",");
				}
			}
		}
		IDataSet data = getDataSet();
		if(this.db2 != null) {
			log.info("Preparing Output Database");
			Connection db1conn = this.db1.getConnection();
			Connection db2conn = this.db2.getConnection();
			Map<Integer,Map<String, String>> inputDbTypes = getDbTypes(db1conn, "input");
			Map<Integer,Map<String, String>> outputDbTypes = getDbTypes(db2conn, "output");
			ResultSet tableData = db2conn.getMetaData().getTables(db2conn.getCatalog(), null, "%", tableTypes);
			while(tableData.next()) {
				String db2tableName = tableData.getString("TABLE_NAME");
				for(String db1table : data.getTableNames()) {
					if(db1table.trim().equalsIgnoreCase(db2tableName.trim())) {
						log.debug("Droping table '"+db2tableName+"' from output database");
						String sql = "DROP TABLE "+db2tableName;
						log.trace("Drop Table SQL Query:\n"+sql);
						db2conn.createStatement().executeUpdate(sql);
					}
				}
			}
			for(String table : data.getTableNames()) {
				// get record set
				log.debug("Creating table '"+table+"' in output database");
				ResultSet columnRS = db1conn.getMetaData().getColumns(null, null, table, null);
				int count = 0;
				StringBuilder createTableSB = new StringBuilder();
				createTableSB.append("CREATE TABLE "+table+" (");
				while(columnRS.next()) {
					if(columnRS.getString("TABLE_NAME").equals(table)) {
						String colName = columnRS.getString("COLUMN_NAME");
						log.debug("Getting column information for '"+colName+"'");
						Integer typeCode = Integer.valueOf(columnRS.getInt("DATA_TYPE"));
						int size = columnRS.getInt("COLUMN_SIZE");
						if(!outputDbTypes.containsKey(typeCode)) {
							if(typeCode.intValue() == Types.BIT) {
								typeCode = Integer.valueOf(Types.BOOLEAN);
								size = 0;
							} else { //TODO: more type conversion issues possible, make this if/else more exhaustive
								if(inputDbTypes.containsKey(typeCode)) {
									log.warn("Output database does not support datatype '"+inputDbTypes.get(typeCode).get("TYPE_NAME")+"': using VARCHAR");
								} else {
									log.error("Unknown datatype code '"+typeCode+"': using VARCHAR");
								}
								typeCode = Integer.valueOf(Types.VARCHAR);
							}
						} else {
//							log.trace("typeCode: "+typeCode);
						}
						Map<String, String> map = outputDbTypes.get(typeCode);
						String typeName = map.get("TYPE_NAME");
						String params = map.get("CREATE_PARAMS");
						if(StringUtils.isBlank(params)) {
							params = map.get("PARAMS");
						}
						boolean needParam = (StringUtils.isNotBlank(params) && (size != 0));
//						log.trace("column '"+colName+"': "+typeCode+" => '"+typeName+((needParam)?"("+size+")":"")+"'");
						if(count != 0) {
							createTableSB.append(',');
						}
						createTableSB.append("\n  ");
						createTableSB.append(colName);
						createTableSB.append(" ");
						createTableSB.append(typeName);
						if(needParam) {
							createTableSB.append("(");
							createTableSB.append(size);
							createTableSB.append(")");
						}
						count++;
					}
				}
				createTableSB.append("\n)");
				log.trace("Create Table SQL Query:\n"+createTableSB);
				db2conn.createStatement().executeUpdate(createTableSB.toString());
			}
			log.info("Dumping Dataset To Output");
			DatabaseOperation.INSERT.execute(this.db2, data);
			log.info("Dataset Output Complete");
		}
		if(this.outFile != null) {
			FlatDtdDataSet.write(data, this.outFile);
		}
	}
	
	/**
	 * Get the dataset
	 * @return the dataset
	 * @throws IOException error writting to flat file
	 * @throws DataSetException error building dataset
	 * @throws SQLException error communicating with database
	 */
	private IDataSet getDataSet() throws DataSetException, IOException, SQLException {
		IDataSet data;
		if(this.db1 != null) {
			DatabaseConfig config = this.db1.getConfig();
			String preEscapePattern = "";
			String postEscaptePattern = "";
			String[] tableTypes = new String[]{"TABLE"};
			if(this.dbUnitFeatures != null) {
				for(String feature : this.dbUnitFeatures.keySet()) {
					Object b = this.dbUnitFeatures.get(feature);
					if(feature.trim().equals("http://www.dbunit.org/properties/escapePattern")) {
						String[] temp = b.toString().split("?", 2);
						preEscapePattern = temp[0];
						if(temp.length > 1) {
							postEscaptePattern = temp[1];
						} else {
							postEscaptePattern = temp[0];
						}
					} else if(feature.trim().equals("http://www.dbunit.org/properties/tableType")) {
						tableTypes = b.toString().split(",");
					}
					log.debug("Setting '"+feature+"' to '"+b+"'");
					config.setProperty(feature, b);
				}
			}
			if(this.tables != null && this.tables.length > 0) {
				// partial database export
				log.info("Constructing Dataset Based on Given Tables");
			} else {
				// full database export
				log.info("Constructing Full Database Dataset");
				Set<String> tableSet = new HashSet<String>();
				ResultSet tableRS = this.db1.getConnection().getMetaData().getTables(null, null, null, tableTypes);
				while(tableRS.next()) {
					tableSet.add(preEscapePattern+tableRS.getString("TABLE_NAME")+postEscaptePattern);
				}
				this.tables = tableSet.toArray(new String[]{});
			}
			QueryDataSet partialDataSet = new QueryDataSet(this.db1);
			for(String table : this.tables) {
				log.debug("Adding table '"+table+"' to dataset");
				String sql = "SELECT * FROM "+table;
				log.trace("Query for table '"+table+"':\n"+sql);
				partialDataSet.addTable(table, sql);
			}
			data = partialDataSet;
		} else if(this.inFile != null) {
			if(this.dbUnitFeatures != null) {
				for(String feature : this.dbUnitFeatures.keySet()) {
					log.warn("feature '"+feature+"' not supported for input files");
				}
			}
			data = new FlatDtdDataSet(this.inFile);
		} else {
			throw new IllegalStateException("inputFile or input database should be initialized!");
		}
		return data;
	}

	/**
	 * Get the database type information for this connection
	 * @param db the database connection
	 * @param dbName the name for this database
	 * @return the type information
	 * @throws SQLException error getting information
	 */
	private Map<Integer, Map<String, String>> getDbTypes(Connection db, String dbName) throws SQLException {
		Map<Integer,Map<String, String>> dbTypes = new HashMap<Integer, Map<String, String>>();
		ResultSet dbTypeInfo = db.getMetaData().getTypeInfo();
		ResultSetMetaData rsmd = dbTypeInfo.getMetaData();
		log.debug("Building type code mappings for "+dbName+" database");
		while(dbTypeInfo.next()) {
			Integer typeCode = Integer.valueOf(dbTypeInfo.getInt("DATA_TYPE"));
//			String typeName = dbTypeInfo.getString("TYPE_NAME");
			Map<String, String> map = null;
			if(!dbTypes.containsKey(typeCode)) {
//				log.trace("Adding mapping information for typecode "+typeCode+":");
				map = new HashMap<String, String>();
			} else {
//				log.trace("Already contained mapping information for typecode: "+typeCode);
			}
			for(int x = 1; x <= rsmd.getColumnCount(); x++) {
				String colVal = dbTypeInfo.getString(x);
				String colName = rsmd.getColumnName(x);
//				log.trace("'"+colName+"' => "+(colVal == null?"":"'")+colVal+(colVal == null?"":"'"));
				if(map != null) {
					map.put(colName, colVal);
				}
			}
			if(map != null) {
				dbTypes.put(typeCode, map);
			}
		}
		return dbTypes;
	}

	/**
	 * Get the ArgParser for this task
	 * @return the ArgParser
	 */
	private static ArgParser getParser() {
		ArgParser parser = new ArgParser("DatabaseClone");
		parser.addArgument(new ArgDef().setLongOpt("inputDriver").withParameter(true, "JDBC_DRIVER").setDescription("jdbc driver class for input database").setRequired(false));
		parser.addArgument(new ArgDef().setLongOpt("inputConnection").withParameter(true, "JDBC_CONN").setDescription("jdbc connection string for input database").setRequired(false));
		parser.addArgument(new ArgDef().setLongOpt("inputUsername").withParameter(true, "USERNAME").setDescription("database username for input database").setRequired(false));
		parser.addArgument(new ArgDef().setLongOpt("inputPassword").withParameter(true, "PASSWORD").setDescription("database password for input database").setRequired(false));
		parser.addArgument(new ArgDef().setLongOpt("inputFile").withParameter(true, "FILE_PATH").setDescription("use this database state file as the input database").setRequired(false));
		parser.addArgument(new ArgDef().setLongOpt("outputDriver").withParameter(true, "JDBC_DRIVER").setDescription("jdbc driver class for output database").setRequired(false));
		parser.addArgument(new ArgDef().setLongOpt("outputConnection").withParameter(true, "JDBC_CONN").setDescription("jdbc connection string for output database").setRequired(false));
		parser.addArgument(new ArgDef().setLongOpt("outputUsername").withParameter(true, "USERNAME").setDescription("database username for output database").setRequired(false));
		parser.addArgument(new ArgDef().setLongOpt("outputPassword").withParameter(true, "PASSWORD").setDescription("database password for output database").setRequired(false));
		parser.addArgument(new ArgDef().setLongOpt("outputFile").withParameter(true, "FILE_PATH").setDescription("output the state of the input database in this file").setRequired(false));
		parser.addArgument(new ArgDef().setShortOption('t').setLongOpt("tableName").withParameters(true, "TABLE_NAME").setDescription("a single database table name [have multiple -t flags, one for each table names]").setRequired(false));
		parser.addArgument(new ArgDef().setLongOpt("validTableType").withParameters(true, "TABLE_TYPE").setDescription("(deprecated) a single table type ('TABLE', 'VIEW', etc) Defaults to just 'TABLE' [have multiple --validTableType flags, one for each table type]").setRequired(false));
		parser.addArgument(new ArgDef().setShortOption('D').setLongOpt("DBUnitFeature").withParameterValueMap("FEATURE", "VALUE").setDescription("Use VALUE for the DBUnit FEATURE (should be (Boolean['true'/'false'] or a comma separated list))").setRequired(false));
		return parser;
	}
	
	/**
	 * Main method
	 * @param args commandline arguments
	 */
	public static void main(String... args) {
		Exception error = null;
		try {
			InitLog.initLogger(args, getParser());
			log.info(getParser().getAppName() + ": Start");
			new DatabaseClone(args).execute();
		} catch(IllegalArgumentException e) {
			log.error(e.getMessage());
			log.debug("Stacktrace:",e);
			System.out.println(getParser().getUsage());
			error = e;
		} catch(UsageException e) {
			log.info("Printing Usage:");
			System.out.println(getParser().getUsage());
			error = e;
		} catch(Exception e) {
			log.error(e.getMessage());
			log.debug("Stacktrace:",e);
			error = e;
		} finally {
			log.info(getParser().getAppName() + ": End");
			if(error != null) {
				System.exit(1);
			}
		}
	}
}