JDBCInstanceWriter.java

/*
 * Copyright (c) 2012 Data Harmonisation Panel
 * 
 * All rights reserved. This program and the accompanying materials are made
 * available under the terms of the GNU Lesser General Public License as
 * published by the Free Software Foundation, either version 3 of the License,
 * or (at your option) any later version.
 * 
 * You should have received a copy of the GNU Lesser General Public License
 * along with this distribution. If not, see <http://www.gnu.org/licenses/>.
 * 
 * Contributors:
 *     HUMBOLDT EU Integrated Project #030962
 *     Data Harmonisation Panel <http://www.dhpanel.eu>
 */

package eu.esdihumboldt.hale.io.jdbc;

import java.io.IOException;
import java.sql.Array;
import java.sql.Connection;
import java.sql.PreparedStatement;
import java.sql.ResultSet;
import java.sql.SQLException;
import java.sql.Statement;
import java.util.ArrayList;
import java.util.Collection;
import java.util.HashMap;
import java.util.HashSet;
import java.util.LinkedHashSet;
import java.util.List;
import java.util.Map;
import java.util.Set;

import javax.xml.namespace.QName;

import de.fhg.igd.slf4jplus.ALogger;
import de.fhg.igd.slf4jplus.ALoggerFactory;
import eu.esdihumboldt.hale.common.core.io.IOProviderConfigurationException;
import eu.esdihumboldt.hale.common.core.io.ProgressIndicator;
import eu.esdihumboldt.hale.common.core.io.report.IOReport;
import eu.esdihumboldt.hale.common.core.io.report.IOReporter;
import eu.esdihumboldt.hale.common.core.io.report.impl.IOMessageImpl;
import eu.esdihumboldt.hale.common.instance.io.impl.AbstractInstanceWriter;
import eu.esdihumboldt.hale.common.instance.model.Instance;
import eu.esdihumboldt.hale.common.instance.model.InstanceCollection;
import eu.esdihumboldt.hale.common.instance.model.ResourceIterator;
import eu.esdihumboldt.hale.common.instance.model.TypeFilter;
import eu.esdihumboldt.hale.common.schema.geometry.GeometryProperty;
import eu.esdihumboldt.hale.common.schema.model.ChildDefinition;
import eu.esdihumboldt.hale.common.schema.model.PropertyDefinition;
import eu.esdihumboldt.hale.common.schema.model.TypeDefinition;
import eu.esdihumboldt.hale.common.schema.model.constraint.property.AutoGenerated;
import eu.esdihumboldt.hale.common.schema.model.constraint.property.NillableFlag;
import eu.esdihumboldt.hale.common.schema.model.constraint.property.Reference;
import eu.esdihumboldt.hale.common.schema.model.constraint.type.GeometryType;
import eu.esdihumboldt.hale.io.jdbc.constraints.DatabaseTable;
import eu.esdihumboldt.hale.io.jdbc.constraints.DefaultValue;
import eu.esdihumboldt.hale.io.jdbc.constraints.SQLArray;
import eu.esdihumboldt.hale.io.jdbc.constraints.SQLType;
import eu.esdihumboldt.hale.io.jdbc.constraints.internal.GeometryAdvisorConstraint;

/**
 * Writes instance to a database through a JDBC connection.
 * 
 * @author Simon Templer
 */
public class JDBCInstanceWriter extends AbstractInstanceWriter
		implements JDBCConstants, JDBCProvider {

	private static final ALogger log = ALoggerFactory.getLogger(JDBCInstanceWriter.class);

	/**
	 * Name of the parameter specifying if the data may be written unordered.
	 */
	private static final String PARAM_UNORDERED = "unordered";

	private final Map<TypeDefinition, Map<Object, Long>> typAuto = new HashMap<>();

	private Map<TypeDefinition, Boolean> visitedType;

	/**
	 * Default constructor.
	 */
	public JDBCInstanceWriter() {
		super();

		addSupportedParameter(PARAM_USER);
		addSupportedParameter(PARAM_PASSWORD);
		addSupportedParameter(PARAM_UNORDERED);
	}

	@Override
	public boolean isCancelable() {
		return true;
	}

	@Override
	public void checkCompatibility() throws IOProviderConfigurationException {
		super.checkCompatibility();

		// XXX check if the target schema is a JDBC schema?
		// this export needs the SQLTypes being set on the property types of
		// instances to write
	}

	@Override
	public Connection getConnection() throws SQLException {
		return JDBCConnection.getConnection(this);
	}

	@Override
	protected IOReport execute(ProgressIndicator progress, IOReporter reporter)
			throws IOProviderConfigurationException, IOException {
		InstanceCollection instances = getInstances();

		Connection connection = null;
		try {
			// connect to the database
			try {
				connection = getConnection();
			} catch (Exception e) {
				reporter.error(new IOMessageImpl(e.getLocalizedMessage(), e));
				reporter.setSuccess(false);
				reporter.setSummary("Failed to connect to database.");
				return reporter;
			}

			if (isWriteUnordered()) {
				// write instances as they come in
				writeInstances(connection, instances, progress, reporter);
			}
			else {
				// write instances based on type order needed for insert
				// (to avoid violating constraints)
				Set<TypeDefinition> sortedSet = getSortedSchemas(
						getTargetSchema().getMappingRelevantTypes());

				for (TypeDefinition td : sortedSet) {
					writeInstances(connection, instances.select(new TypeFilter(td)), progress,
							reporter);
				}
			}

			reporter.setSuccess(true);
		} catch (Exception e) {
			reporter.error(new IOMessageImpl(e.getLocalizedMessage(), e));
			reporter.setSuccess(false);
			reporter.setSummary("Saving instances to database failed.");
		} finally {
			if (connection != null) {
				try {
					connection.close();
				} catch (SQLException e) {
					// ignore
				}
			}
			progress.end();
		}

		return reporter;
	}

	@Override
	public boolean isPassthrough() {
		return isWriteUnordered();
	}

	/**
	 * @return if instances are written unordered
	 */
	protected boolean isWriteUnordered() {
		return getParameter(PARAM_UNORDERED).as(Boolean.class, false);
	}

	/**
	 * Write instances to a database connection Auto incremental keys are
	 * written to the database with the new generated values, and updated the
	 * foreign key values with this newly generated auto incremental values of
	 * the referenced column.
	 * 
	 * <p>
	 * <b>Limitation:</b>
	 * </p>
	 * This implementation for auto generated key insert is supported only for
	 * those tables which has only single auto generated fields.
	 * 
	 * @param connection the database connection
	 * @param instances the instances to write
	 * @param progress the progress indicator
	 * @param reporter the reporter
	 * @throws Exception if saving the instances fails
	 */
	private void writeInstances(Connection connection, InstanceCollection instances,
			ProgressIndicator progress, IOReporter reporter) throws Exception {
		connection.setAutoCommit(false);

		boolean trackProgress = instances.hasSize();
		progress.begin("Write instances to database",
				(trackProgress) ? (instances.size()) : (ProgressIndicator.UNKNOWN));

		// maps type definitions to prepared statements
		Map<TypeDefinition, Map<Set<QName>, PreparedStatement>> typeStatements = new HashMap<TypeDefinition, Map<Set<QName>, PreparedStatement>>();
		Map<TypeDefinition, Map<Set<QName>, Integer>> typeCount = new HashMap<TypeDefinition, Map<Set<QName>, Integer>>();

		ResourceIterator<Instance> it = instances.iterator();
		try {
			while (it.hasNext() && !progress.isCanceled()) {
				Instance instance = it.next();
				TypeDefinition type = instance.getDefinition();

				Set<QName> properties = new HashSet<QName>();
				for (QName property : instance.getPropertyNames())
					properties.add(property);
				filterInsertProperties(type, properties);

				// per type count
				Map<Set<QName>, Integer> typeCountMap = typeCount.get(type);
				if (typeCountMap == null) {
					typeCountMap = new HashMap<Set<QName>, Integer>();
					typeCount.put(type, typeCountMap);
				}
				Integer count = typeCountMap.get(properties);
				if (count == null)
					count = 0;
				typeCountMap.put(properties, count + 1);

				// get prepared statement for instance type
				PreparedStatement statement = getInsertStatement(type, properties, typeStatements,
						connection);

				// populate insert statement with values
				populateInsertStatementOrExecuteAutoIncStatement(statement, properties, instance,
						reporter, connection);

				if (count % 100 == 0) {
					statement.executeBatch();
					// TODO statement.getGeneratedKeys() / does not work with
					// batches for PostgreSQL
				}

				if (trackProgress) {
					progress.advance(1);
				}
			}

			// execute remaining batches
			for (Map<Set<QName>, PreparedStatement> typeSpecificMap : typeStatements.values()) {
				if (progress.isCanceled())
					break;
				for (PreparedStatement statement : typeSpecificMap.values()) {
					if (progress.isCanceled())
						break;
					statement.executeBatch();
				}
			}

			if (!progress.isCanceled())
				connection.commit();
			else
				connection.rollback();
		} catch (Exception e) {
			if (e instanceof SQLException) {
				SQLException next = ((SQLException) e).getNextException();
				while (next != null) {
					log.error("SQL exception", next);
					next = next.getNextException();
				}
			}
			try {
				connection.rollback();
			} catch (Exception e1) {
				// ignore
			}
			throw e;
		} finally {
			// close iterator
			it.close();

			// close statements
			for (Map<Set<QName>, PreparedStatement> typeSpecificMap : typeStatements.values())
				for (PreparedStatement statement : typeSpecificMap.values())
					statement.close();
		}

		// right now cancel => rollback. Otherwise this would have to be in
		// front of it.close()...
//		if (progress.isCanceled() && it.hasNext()) {
//			reporter.error(new IOMessageImpl("Saving to database was canceled, not all instances were saved.", null));
//		}
	}

	/**
	 * 
	 * This helps to store the mapping of the old value to the auto generated
	 * value for a particular type. If same type appears twice that means the
	 * sorting and filtering of the types has not been done properly.
	 * 
	 * @param type a type definition whose child has the auto incremented
	 *            constraint
	 * @param orgId original id as string
	 * @param genId auto incremented id
	 */
	private void addIDMapping(TypeDefinition type, Object orgId, Long genId) {
		orgId = processLookupId(orgId);

		Map<Object, Long> autoInc = typAuto.get(type);
		if (autoInc == null) {
			autoInc = new HashMap<>();
			typAuto.put(type, autoInc);
		}
		autoInc.put(orgId, genId);
	}

	private Object processLookupId(Object orgId) {
		if (orgId instanceof Number) {
			return ((Number) orgId).longValue();
		}
		return orgId;
	}

	/**
	 * 
	 * Retrieves the target schema and orders it based on the references, so
	 * that there should not be any integrity constrain violation exception.
	 * <p>
	 * <b> Algorithm for sorting </b>
	 * </p>
	 * The steps for sorting the types based on references are as follows:
	 * <ol>
	 * <li>Create a new <code>set</code> for the sorted types</li>
	 * <li>Create a visited type map to keep hold of the types whose references
	 * have already been checked.
	 * 
	 * <li>Iterate over the types that needs to be sorted
	 * <ol>
	 * <li>If the type is already in the visited map means references for this
	 * type have already been checked. (Again checking could lead to cyclic
	 * referencing in the infinite loop)</li>
	 * <li>Exit the current recursion.</li>
	 * <li>Iterate over the child definition of the types and check for the
	 * referenced types</li>
	 * <li>If there are reference types then iterate over each referenced type
	 * and recurse to the step (3.1) with this iterated referenced type</li>
	 * <li>Traverse till the type whose childs do not have any reference, add
	 * this type to the sorted set first</li>
	 * <li>Add the iterated type</li>
	 * 
	 * </ol>
	 * </li>
	 * <li>Add the sorted types in the sorted set</li>
	 * </ol>
	 * 
	 * 
	 * @param collection types to be sorted
	 * 
	 * @return returns the sorted schema based on the references
	 * 
	 */
	public Set<TypeDefinition> getSortedSchemas(Collection<? extends TypeDefinition> collection) {

		Set<TypeDefinition> sortedSet = new LinkedHashSet<TypeDefinition>();
		visitedType = new HashMap<TypeDefinition, Boolean>();
		for (TypeDefinition td : collection) {
			if (!sortedSet.contains(td)) {
				List<TypeDefinition> references = getReferencesTd(td);
				if (references != null)
					sortedSet.addAll(references);
			}

		}

		return sortedSet;
	}

	/**
	 * gets the list of the referenced type definitions. This is used for
	 * sorting the types based on the references. This method recursively calls
	 * itself till it finds a type which does not have any reference and that
	 * type to the sorted list as first element.
	 * 
	 * @param td type definition of the reference
	 * @return list of type definitions of referenced
	 */
	private List<TypeDefinition> getReferencesTd(TypeDefinition td) {

		if (visitedType.containsKey(td) && visitedType.get(td)) {
			return null;
		}
		else {
			visitedType.put(td, true);
		}
		List<TypeDefinition> referencedTD = new ArrayList<TypeDefinition>();
		for (ChildDefinition<?> cd : td.getChildren()) {
			PropertyDefinition pd = cd.asProperty();
			if (pd != null) {
				Reference ref = pd.getConstraint(Reference.class);
				// check if it has the references
				if (ref.getReferencedTypes() != null) {
					for (TypeDefinition t : ref.getReferencedTypes()) {
						List<TypeDefinition> references = getReferencesTd(t);
						if (references != null)
							referencedTD.addAll(references);
					}
				}
			}
		}
		if (!referencedTD.contains(td))
			referencedTD.add(td);

		return referencedTD;
	}

	/**
	 * Filters the set of properties to only contain properties that can be used
	 * for inserting (e. g. no groups).
	 * 
	 * @param type the type definition
	 * @param properties the available properties
	 */
	private void filterInsertProperties(TypeDefinition type, Set<QName> properties) {
		for (ChildDefinition<?> child : type.getChildren()) {
			if (properties.contains(child.getName())) {
				if (child.asProperty() == null) {
					// remove it since it is a group and no property
					// XXX warn?
					properties.remove(child.getName());
				}
			}
		}
	}

	/**
	 * Create a prepared insert statement, based on the given type definition.
	 * Currently, auto incremented fields are not inserted into the statement.
	 * 
	 * @param type the type definition
	 * @param properties the set properties of the instance for which this
	 *            statement is
	 * @param typeStatements the already created statements
	 * @param connection the database connection
	 * @return the insert statement
	 * @throws SQLException if creating the prepared statement fails
	 */
	private PreparedStatement getInsertStatement(TypeDefinition type, Set<QName> properties,
			Map<TypeDefinition, Map<Set<QName>, PreparedStatement>> typeStatements,
			Connection connection) throws SQLException {
		Map<Set<QName>, PreparedStatement> typeSpecificMap = typeStatements.get(type);
		if (typeSpecificMap == null) {
			typeSpecificMap = new HashMap<Set<QName>, PreparedStatement>();
			typeStatements.put(type, typeSpecificMap);
		}

		PreparedStatement result = typeSpecificMap.get(properties);

		if (result == null) {
			String tableName = type.getConstraint(DatabaseTable.class).getFullTableName();

			// create prepared statement SQL
			StringBuffer pSql = new StringBuffer();

			pSql.append("INSERT INTO ");
			pSql.append(tableName);
			pSql.append(" (");

			StringBuffer valuesSql = new StringBuffer();

			boolean first = true;
			for (QName property : properties) {

				AutoGenerated auto = ((PropertyDefinition) type.getChild(property))
						.getConstraint(AutoGenerated.class);
				// Auto incremented fields are not included in the insert
				// statement (if writing is ordered)
				if (!auto.isEnabled() || isWriteUnordered()) {
					if (first)
						first = false;
					else {
						pSql.append(", ");
						valuesSql.append(",");
					}

					pSql.append('"').append(property.getLocalPart()).append('"');
					valuesSql.append('?');
				}
			}
			pSql.append(") VALUES (");
			pSql.append(valuesSql);
			pSql.append(")");

			// XXX Actually we don't necessarily need the auto generated keys,
			// we need the primary key!
			// XXX , Statement.RETURN_GENERATED_KEYS does not work with batches
			// in PostgreSQL
			// Auto generated keys are returned, for creating a map with old
			// value and auto generated.
			// As it does not work with batches in PostGreSQL, each instance
			// which has auto incremented ids are written one at a time and
			// returned the auto generated key.
			result = connection.prepareStatement(pSql.toString(), Statement.RETURN_GENERATED_KEYS);
			typeSpecificMap.put(properties, result);
		}

		return result;
	}

	/**
	 * Populate a prepared insert statement with values from the given instance.
	 * Checks if the property has auto incremental constraint, if it has then it
	 * maps the original old id to the id that is auto incremented while
	 * inserting the value. This mapping is used when inserting the foreign key
	 * values associated with those auto incremented column ids whose value has
	 * been changed. Thus, insertion of the foreign key wont fail. It will
	 * either execute the statement directly or add it into the batches
	 * depending upon the auto incremented flag.
	 * 
	 * 
	 * @param statement the insert statement
	 * @param properties the properties to fill the statement with
	 * @param instance the instance
	 * @param reporter the reporter
	 * @param conn Connection (used for Geometry conversion for oracle)
	 * @throws SQLException if configuring the statement fails
	 */
	private void populateInsertStatementOrExecuteAutoIncStatement(PreparedStatement statement,
			Set<QName> properties, Instance instance, IOReporter reporter, Connection conn)
					throws SQLException {
		TypeDefinition type = instance.getDefinition();

		int index = 1;
		Object oldValue = null;
		boolean isAutoIncremented = false;
		for (QName propertyName : properties) {
			PropertyDefinition property = (PropertyDefinition) type.getChild(propertyName);
			Object[] values = instance.getProperty(propertyName);
			SQLType sqlType = property.getPropertyType().getConstraint(SQLType.class);
			if (!sqlType.isSet()) {
				reporter.error(new IOMessageImpl(
						"SQL type not set. Please only export to schemas read from a database.",
						null));
				statement.setObject(index, null);
				continue;
			}

			SQLArray arrayInfo = property.getPropertyType().getConstraint(SQLArray.class);

			Object value;
			if (arrayInfo.isArray() && arrayInfo.getDimension() <= 1) {
				// array as multiple occurrence property
				value = (values == null) ? (new Object[0]) : (values);
			}
			else {
				// single value
				if (values != null && values.length > 1)
					reporter.warn(new IOMessageImpl(
							"Multiple values for a property. Only exporting first.", null));
				value = (values == null || values.length == 0) ? null : values[0];
			}

			AutoGenerated auto = property.getConstraint(AutoGenerated.class);
			if (!isWriteUnordered()) {
				// only able to handle auto-incremented IDs if instances are
				// written ordered

				if (auto.isEnabled()) {
					isAutoIncremented = true;
					if (value != null) {
						oldValue = value;
					}
					continue;
				}
			}

			Reference ref = property.getConstraint(Reference.class);
			if (ref.getReferencedTypes() != null) {
				TypeDefinition td = (TypeDefinition) ref.getReferencedTypes().toArray()[0];
				Map<Object, Long> marshMallow = typAuto.get(td);
				if (marshMallow != null && value != null) {
					// lookup identifier for reference
					value = marshMallow.get(processLookupId(value));
				}
			}

			if (values == null || values.length == 0) {
				// XXX The default value could be a function call.
				// Better would be to leave the column out of the insert
				// statement, or set it to the SQL keyword "DEFAULT".
				DefaultValue defaultValue = property.getConstraint(DefaultValue.class);
				if (defaultValue.isSet())
					statement.setObject(index, defaultValue.getValue(), sqlType.getType());
				else if (property.getConstraint(NillableFlag.class).isEnabled())
					statement.setNull(index, sqlType.getType());
				else {
					// no default, not nillable, will not work...
					// set it to null here and let query fail (probably)
					// XXX maybe skip this insert?
					statement.setNull(index, sqlType.getType());
					reporter.warn(new IOMessageImpl(
							"Property no value, not nillable, no default value, insert will probably fail.",
							null));
				}
			}
			else if (value == null)
				statement.setNull(index, sqlType.getType());

			else
				setStatementParameter(statement, index, value, property, sqlType.getType(),
						reporter, conn);

			index++;
		}
		if (isAutoIncremented) {
			statement.execute();
			ResultSet rs = statement.getGeneratedKeys();
			Long newValue = null;
			while (rs.next()) {
				newValue = rs.getLong(1);
			}

			addIDMapping(type, oldValue, newValue);
		}
		else {
			statement.addBatch();
		}
	}

	/**
	 * Set a prepared statement parameter value.
	 * 
	 * @param statement the prepared statement
	 * @param index the parameter index
	 * @param value the value, not <code>null</code>
	 * @param propertyDef the associated property definition
	 * @param sqlType the SQL type
	 * @param reporter the reporter
	 * @param conn Connection (currently used for geometry conversion for oracle
	 *            database
	 * @throws SQLException if setting the parameter fails
	 */
	@SuppressWarnings("unchecked")
	private void setStatementParameter(PreparedStatement statement, int index, Object value,
			PropertyDefinition propertyDef, int sqlType, IOReporter reporter, Connection conn)
					throws SQLException {
		if (propertyDef.getPropertyType().getConstraint(GeometryType.class).isGeometry()) {
			// is a geometry column

			// get the geometry advisor
			@SuppressWarnings("rawtypes")
			GeometryAdvisor advisor = propertyDef.getPropertyType()
					.getConstraint(GeometryAdvisorConstraint.class).getAdvisor();
			if (advisor != null) {
				// use the advisor to convert the geometry
				if (value instanceof GeometryProperty<?>) {
					// XXX JTS geometry conversion needed beforehand?
					try {
						value = advisor.convertGeometry((GeometryProperty<?>) value,
								propertyDef.getPropertyType(), conn);
					} catch (Exception e) {
						reporter.error(
								new IOMessageImpl("Something went wrong during conversion", e));
					}
				}
				else {
					reporter.error(new IOMessageImpl(
							"Geometry value is not of type GeometryProperty and could thus not be converted for the database",
							null));
				}
			}
		}

		SQLArray arrayInfo = propertyDef.getPropertyType().getConstraint(SQLArray.class);
		if (arrayInfo.isArray()) {
			// is an array column

			Object[] values;
			if (value.getClass().isArray()) {
				values = (Object[]) value;
			}
			else {
				values = new Object[] { value };
			}

			// FIXME for multi-dimensional arrays, make sure internal structures
			// are arrays?

			// use SQL array as value
			Array array = conn.createArrayOf(arrayInfo.getElementTypeName(), values);
			value = array;

			// FIXME collect arrays to allow them to be freed after the
			// statement is executed?
		}

		// TODO handling of other types?

		// set the value
		statement.setObject(index, value, sqlType);
	}

	@Override
	protected String getDefaultTypeName() {
		return "Database";
	}
}