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

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

import javax.inject.Inject;
import javax.inject.Named;

import net.sourceforge.jtds.jdbc.Driver;
import org.etlunit.Log;
import org.etlunit.RuntimeSupport;
import org.etlunit.TestExecutionError;
import org.etlunit.feature.database.*;
import org.etlunit.parser.ETLTestValueObject;
import org.etlunit.util.IOUtils;
import org.etlunit.util.VelocityUtil;

import java.io.File;
import java.io.IOException;
import java.net.URL;
import java.sql.Connection;
import java.sql.SQLException;
import java.sql.Statement;
import java.util.HashMap;
import java.util.Map;

public class SqlServerDatabaseImplementation extends BaseDatabaseImplemenation
{
	private Log applicationLog;
	private RuntimeSupport runtimeSupport;
	private DatabaseFeatureModule databaseFeatureModule;
	private static final int MASTER = -1;
	private static final int DATABASE = 0;

	@Inject
	public void setRuntimeSupport(RuntimeSupport runtimeSupport)
	{
		this.runtimeSupport = runtimeSupport;
	}

	@Inject
	public void setDatabaseFeatureModule(DatabaseFeatureModule databaseFeatureModule)
	{
		this.databaseFeatureModule = databaseFeatureModule;
	}

	private Map identityStatus = new HashMap();

	@Override
	public void prepareConnectionForInsert(Connection connection, RelationalDataSet target, String mode) throws Exception
	{
		if (mode == null)
		{
			mode = "[[DEFAULT]]";
		}

		if (target.hasIdentityColumns())
		{
			applicationLog.info("Enabling identity insert for " + target.getQualifiedName());

			RelationalDataSet currIdent = identityStatus.get(mode);
			Statement st = connection.createStatement();

			if (currIdent != null)
			{
				String sql = "SET IDENTITY_INSERT " + currIdent.getQualifiedName() + " OFF";

				try
				{
					st.execute(sql);
				}
				catch (SQLException exc)
				{
					applicationLog.severe("Error disabling identity insert for " + currIdent.getQualifiedName() + " using sql: " + sql, exc);
				}
			}

			String sql = "SET IDENTITY_INSERT " + target.getQualifiedName() + " ON";

			try
			{
				st.execute(sql);
				identityStatus.put(mode, target);
			}
			catch (SQLException exc)
			{
				applicationLog.severe("Error enabling identity insert for " + target.getQualifiedName() + " using sql: " + sql, exc);
			}
			finally
			{
				st.close();
			}
		}
		else
		{
			applicationLog.info("Identity insert not required for " + target.getQualifiedName());
		}
	}

	@Inject
	public void setLogger(@Named("applicationLog") Log log)
	{
		this.applicationLog = log;
	}

	public data_format getDataFormatForExtract(ETLTestValueObject operands)
	{
		return data_format.delimited;
	}

	public String getImplementationId()
	{
		return "sql-server";
	}

	public Object processOperation(operation op, OperationRequest request) throws UnsupportedOperationException
	{
		switch (op)
		{
			case prepareDatabase:
				doPrepareDatabase(request.getPrepareRequest());
				return null;

			default:
				throw new UnsupportedOperationException();
			case createDatabase:
				doCreateDatabase(request.getInitializeRequest());
				return null;
			case completeDatabaseInitialization:
				doCompleteDatabase(request.getInitializeRequest());
				return null;
		}
	}

	public String getDefaultSchema(DatabaseConnection dc, String mode)
	{
		return "DBO";
	}

	public database_state getDatabaseState(DatabaseConnection dc, String mode)
	{
		return database_state.pass;
	}

	private void doCompleteDatabase(InitializeRequest initializeRequest)
	{
		URL url_materialize_views = getClass().getResource("/materialize_views.vm");

		Map map = new HashMap();
		DatabaseConnection databaseConnection = initializeRequest.getConnection();
		String mode = initializeRequest.getMode();

		try
		{
			// materialize the views into tables
			String materializeText = IOUtils.readURLToString(url_materialize_views);

			String script = VelocityUtil.writeTemplate(materializeText, map);

			String scriptName = "Materialize_Views_" + databaseConnection.getDatabaseName(mode) + ".sql";

			File scriptFile = runtimeSupport.createGeneratedSourceFile(getImplementationId(), scriptName);

			IOUtils.writeBufferToFile(scriptFile, new StringBuffer(script));

			executeScript(script, databaseConnection, mode, false);
		}
		catch (Exception e)
		{
			// TODO:  Get the appropriate exceptions added to this interface method
			throw new IllegalArgumentException(e);
		}
	}

	private void executeScript(String script, DatabaseConnection databaseConnection, String mode, boolean master) throws IOException, TestExecutionError
	{
		final SQLAggregator sqlagg = databaseFeatureModule.resolveSQLRef(script);

		jdbcClient.useStatement(databaseConnection, mode, new JDBCClient.StatementClient()
		{
			public void connection(Connection conn, Statement st, DatabaseConnection connection, String mode, int id) throws Exception
			{
				SQLAggregator.Aggregator statementAggregator = sqlagg.getStatementAggregator();
				while (statementAggregator.hasNext())
				{
					st.execute(statementAggregator.next().getLine());
				}
			}
		}, master ? MASTER : DATABASE);
	}

	private void doCreateDatabase(InitializeRequest initializeRequest)
	{
		URL url_kill = getClass().getResource("/kill_db.vm");
		URL url_create = getClass().getResource("/create_db.vm");

		try
		{
			String killText = IOUtils.readURLToString(url_kill);

			Map map = new HashMap();
			DatabaseConnection databaseConnection = initializeRequest.getConnection();
			String mode = initializeRequest.getMode();
			map.put("databaseLogin", databaseConnection.getLoginName(mode));
			map.put("databaseName", databaseConnection.getDatabaseName(mode));
			map.put("databaseUsername", databaseConnection.getLoginName(mode));
			map.put("databasePassword", databaseConnection.getPassword(mode));

			String script = VelocityUtil.writeTemplate(killText, map);

			String scriptName = "Kill_Database_" + databaseConnection.getDatabaseName(mode) + ".sql";

			File scriptFile = runtimeSupport.createGeneratedSourceFile(getImplementationId(), scriptName);

			IOUtils.writeBufferToFile(scriptFile, new StringBuffer(script));

			// execute the script through sqlcmd
			executeScript(script, databaseConnection, mode, true);

			String createText = IOUtils.readURLToString(url_create);

			script = VelocityUtil.writeTemplate(createText, map);

			scriptName = "Create_Database_" + databaseConnection.getDatabaseName(mode) + ".sql";

			scriptFile = runtimeSupport.createGeneratedSourceFile(getImplementationId(), scriptName);

			IOUtils.writeBufferToFile(scriptFile, new StringBuffer(script));

			executeScript(script, databaseConnection, mode, true);
		}
		catch (Exception e)
		{
			// TODO:  Get the appropriate exceptions added to this interface method
			throw new IllegalArgumentException(e);
		}
	}

	private void doPrepareDatabase(PrepareRequest prepareRequest)
	{
		URL url = getClass().getResource("/purge_db.vm");

		try
		{
			String purgeText = IOUtils.readURLToString(url);
			DatabaseConnection databaseConnection = prepareRequest.getConnection();
			String mode = prepareRequest.getMode();
			String databaseName = databaseConnection.getDatabaseName(mode);

			String scriptName = "Purge_Database_" + databaseConnection.getDatabaseName(mode) + ".sql";

			File scriptFile = runtimeSupport.createGeneratedSourceFile(getImplementationId(), scriptName);

			IOUtils.writeBufferToFile(scriptFile, new StringBuffer(purgeText));

			// execute the script through sqlcmd
			executeScript(purgeText, databaseConnection, mode, false);
		}
		catch (Exception e)
		{
			// TODO:  Get the appropriate exceptions added to this interface method
			throw new IllegalStateException("", e);
		}
	}

	@Override
	protected String getJdbcUrl(DatabaseConnection dc, String mode, int id)
	{
		if (id == MASTER)
		{
			return "jdbc:jtds:sqlserver://" + dc.getServerName() + "/master";
		}
		else
		{
			return "jdbc:jtds:sqlserver://" + dc.getServerName() + "/" + dc.getDatabaseName(mode);
		}
	}

	@Override
	protected String getPassword(DatabaseConnection dc, String mode, int id)
	{
		return id == MASTER ? dc.getAdminPassword() : dc.getPassword(mode);
	}

	@Override
	protected Class getJdbcDriverClass()
	{
		return Driver.class;
	}

	@Override
	protected String getLoginName(DatabaseConnection dc, String mode, int id)
	{
		return id == MASTER ? dc.getAdminUserName() : dc.getLoginName(mode);
	}
}