org.etlunit.feature.sql_server_database.BcpCmd Maven / Gradle / Ivy

Go to download
package org.etlunit.feature.sql_server_database;

import org.etlunit.*;
import org.etlunit.feature.database.DatabaseImplementation;
import org.etlunit.parser.UnEscapist;
import org.etlunit.util.IOUtils;
import org.etlunit.util.StringUtils;
import org.etlunit.util.VelocityUtil;
import org.etlunit.util.regexp.ColumnListExpression;
import org.etlunit.util.regexp.FormatFileColumnExpression;

import java.io.File;
import java.io.IOException;
import java.net.URL;
import java.nio.charset.Charset;
import java.util.*;

public class BcpCmd {
	public enum bcp_direction
	{
		in,
		out,
		queryout
	}

	private final File bcpCmd;
	private final String serverName;
	private final String databaseName;
	private final RuntimeSupport runtimeSupport;
	private final SqlCmd sqlCmd;
	private final String columnDelimiter;
	private final String rowDelimiter;

	public BcpCmd(
		File sqlBinDir,
		RuntimeSupport runtimeSupport,
		String serverName,
		String databaseName,
		String delimiter,
		String rowDelimiter
	) {
		this.bcpCmd = new File(sqlBinDir, "bcp.exe");
		this.sqlCmd = new SqlCmd(sqlBinDir, runtimeSupport);
		this.runtimeSupport = runtimeSupport;
		this.databaseName = databaseName;
		this.serverName = serverName;
		this.columnDelimiter = UnEscapist.escape(delimiter);
		this.rowDelimiter = UnEscapist.escape(rowDelimiter);

		if (!bcpCmd.exists())
		{
			throw new IllegalArgumentException("bcp.exe not found");
		}
	}

	public void importCsvData(
		File dataFile,
		File formatFile,
		String targetTableName,
		String targetTableSchema
	) throws IOException {
		generateFormatFile(targetTableName, targetTableSchema, formatFile);

		ProcessDescription pb = new ProcessDescription(
			bcpCmd.getAbsolutePath())
			.argument(databaseName + "." + targetTableSchema + "." + targetTableName)
			.argument("in")
			.argument(dataFile.getAbsolutePath())
			.argument("-n")
			.argument("-T")
			.argument("-S")
			.argument(serverName)
			.argument("-f")
			.argument(formatFile.getAbsolutePath())
			.argument("-F")
			.argument("2")
			.argument("-c")
			.argument("-t")
			.argument(columnDelimiter)
			.argument("-r")
			.argument(rowDelimiter);

		ProcessFacade process = runtimeSupport.execute(pb);
		process.waitForCompletion();

		int res = process.getCompletionCode();

		String buffered = process.getInputBuffered().toString();

		if (res != 0)
		{
			throw new IOException("Could not load data: " + dataFile.getAbsolutePath() + " : " + process.getInputBuffered());
		}
	}

	public List exportSql(String sql, File targetFile) throws Exception {
		String name = targetFile.getName().replace('.', '_');

		// create a stored procedure because bcp doesn't offer a database switch in every version
		Map bean = new HashMap();

		bean.put("databaseName", databaseName);
		bean.put("procedureName", name);
		bean.put("sql", sql);

		String str = getResourceTemplate("proc_frag", bean);

		File tempScript = runtimeSupport.createGeneratedSourceFile("sql-server", name + "_proc.sql");

		IOUtils.writeBufferToFile(tempScript, new StringBuffer(str));

		sqlCmd.executeScript(serverName, databaseName, tempScript);

		ProcessDescription pb = new ProcessDescription(
			bcpCmd.getAbsolutePath())
			.argument("exec " + databaseName + ".dbo." + name)
			.argument("queryout")
			.argument(targetFile.getAbsolutePath())
			.argument("-n")
			.argument("-T")
			.argument("-S")
			.argument(serverName)
			.argument("-V")
			.argument("90");

		ProcessFacade process = runtimeSupport.execute(pb);
		process.waitForCompletion();

		int res = process.getCompletionCode();

		if (res != 0)
		{
			throw new IOException("Could not execute script: " + targetFile.getAbsolutePath());
		}

		// create an sp to export the columns from the sql
		name = name + "_col_export";

		bean.put("tempTableName", name);

		str = getResourceTemplate("export_sql_columns", bean);

		tempScript = runtimeSupport.createGeneratedSourceFile("sql-server", name + "_proc.sql");
		File tempColumns = runtimeSupport.createGeneratedSourceFile("sql-server", name + ".columns");

		IOUtils.writeBufferToFile(tempScript, new StringBuffer(str));

		// load the procedure into the database
		sqlCmd.executeScript(serverName, databaseName, tempScript);

		// execute the procedure into an export file
		pb = new ProcessDescription(
			bcpCmd.getAbsolutePath())
			.argument("exec " + databaseName + ".dbo." + name)
			.argument("queryout")
			.argument(tempColumns.getAbsolutePath())
			.argument("-n")
			.argument("-T")
			.argument("-S")
			.argument(serverName)
			.argument("-V")
			.argument("90");

		process = runtimeSupport.execute(pb);
		process.waitForCompletion();

		res = process.getCompletionCode();

		if (res != 0)
		{
			throw new IOException("Could not execute export columns script: " + targetFile.getAbsolutePath());
		}

		// load the export file and read the rows into columns
		List colList = new ArrayList();

		String cols = IOUtils.readFileToString(tempColumns);

		ColumnListExpression cle = new ColumnListExpression(cols);

		while (cle.hasNext())
		{
			colList.add(cle.getColumnName());
		}

		return colList;
	}

	private String getResourceTemplate(String path, Map bean) throws Exception {
		URL template = getClass().getResource("/" + path + ".vm");

		if (template == null)
		{
			throw new IllegalStateException("Could not find " + path + ".vm template");
		}

		return VelocityUtil.writeTemplate(IOUtils.readURLToString(template), bean);
	}

	public void generateFormatFile(
			String tableName,
			String tableSchema,
			File format
	) throws IOException {
		ProcessDescription pb = new ProcessDescription(
				bcpCmd.getAbsolutePath())
				.argument(databaseName + "." + tableSchema + "." + tableName)
				.argument("format")
				.argument("nul")
				.argument("-f")
				.argument(format.getAbsolutePath())
				.argument("-T")
				.argument("-S")
				.argument(serverName)
				.argument("-V")
				.argument("90")
				.argument("-x")
				.argument("-w")
				.argument("-t")
				.argument(columnDelimiter)
				.argument("-r")
				.argument(rowDelimiter);

		ProcessFacade process = runtimeSupport.execute(pb);
		process.waitForCompletion();

		int res = process.getCompletionCode();

		if (res != 0)
		{
			throw new IOException("Could not create format file: " + process.getInputBuffered());
		}
	}

	public List exportTableToCsv(
			String tableName,
			String tableSchema,
			List columns,
			DatabaseImplementation.ExtractRequest.column_list_mode columnListMode,
			File data,
			File format
	) throws Exception {
		generateFormatFile(tableName, tableSchema, format);

		String ff = readBcpEncodedFileToString(format);

		List ffColumns = new ArrayList();
		FormatFileColumnExpression ffce = new FormatFileColumnExpression(ff);

		while (ffce.hasNext())
		{
			ffColumns.add(ffce.getColumnName());
		}

		if (columnListMode != DatabaseImplementation.ExtractRequest.column_list_mode.none)
		{
			if (columns == null)
			{
				throw new IllegalStateException("column-list-mode requires columns");
			}

			switch (columnListMode)
			{
				case excludes:
				{
					for (String col : columns)
					{
						if (ffColumns.contains(col))
						{
							ffColumns.remove(col);
						}
						else
						{
							throw new TestExecutionError("Column '" + col + "' in exclude list does not exist in the source table: " + ffColumns);
						}
					}

					break;
				}
				case includes:
				{
					List newCols = new ArrayList();

					for (String col : columns)
					{
						if (ffColumns.contains(col))
						{
							newCols.add(col);
						}
						else
						{
							throw new TestExecutionError("Column '" + col + "' in include list does not exist in the source table: " + ffColumns);
						}
					}

					ffColumns = newCols;
				}

				break;
			}
		}

		StringBuffer select = new StringBuffer("\nSELECT\n");
		StringBuffer order = new StringBuffer("\nFROM\n\t" + tableSchema + "." + tableName + "\nORDER BY\n");

		boolean first = true;

		for (String col : ffColumns)
		{
			if (!first)
			{
				select.append(",\n");
				order.append(",\n");
			}
			else
			{
				first = false;
			}

			select.append("\t");
			select.append(col);
			order.append("\t");
			order.append(col);
		}

		Map bean = new HashMap();

		String procName = StringUtils.sanitize(data.getName(), '_');

		bean.put("databaseName", databaseName);
		bean.put("procedureName", procName);
		bean.put("sql", select.toString() + order.toString() + "\n");

		String str = getResourceTemplate("proc_frag", bean);

		File tempScript = runtimeSupport.createGeneratedSourceFile("sql-server", procName + "_proc.sql");

		IOUtils.writeBufferToFile(tempScript, new StringBuffer(str));

		sqlCmd.executeScript(serverName, databaseName, tempScript);

		ProcessDescription pb = new ProcessDescription(
			bcpCmd.getAbsolutePath())
		.argument("exec " + databaseName + "." + tableSchema + "." + procName)
			.argument("queryout")
			.argument(data.getAbsolutePath())
			.argument("-T")
			.argument("-S")
			.argument(serverName)
			.argument("-V")
			.argument("90")
			.argument("-c")
			.argument("-t")
			.argument(columnDelimiter)
			.argument("-r")
			.argument(rowDelimiter);

		ProcessFacade process = runtimeSupport.execute(pb);
		process.waitForCompletion();

		int res = process.getCompletionCode();

		if (res != 0)
		{
			throw new IOException("Could not execute bcp command: " + process.getInputBuffered());
		}

		return ffColumns;
	}

	static String readBcpEncodedFileToString(File file) throws IOException
	{
		return new String(IOUtils.readFileToString(file).getBytes(), Charset.forName("UTF-16LE"));
	}
}