RelationalEntityFinder.java

package org.codefilarete.stalactite.engine.runtime;

import java.sql.ResultSet;
import java.util.Collections;
import java.util.HashMap;
import java.util.HashSet;
import java.util.Iterator;
import java.util.Map;
import java.util.Set;
import java.util.function.Consumer;
import java.util.function.Function;
import java.util.stream.Collectors;

import org.codefilarete.stalactite.engine.SelectExecutor;
import org.codefilarete.stalactite.engine.configurer.builder.BuildLifeCycleListener;
import org.codefilarete.stalactite.engine.configurer.builder.PersisterBuilderContext;
import org.codefilarete.stalactite.engine.runtime.load.EntityInflater;
import org.codefilarete.stalactite.engine.runtime.load.EntityJoinTree;
import org.codefilarete.stalactite.engine.runtime.load.EntityTreeInflater;
import org.codefilarete.stalactite.engine.runtime.load.EntityTreeQueryBuilder;
import org.codefilarete.stalactite.engine.runtime.load.EntityTreeQueryBuilder.EntityTreeQuery;
import org.codefilarete.stalactite.engine.runtime.query.EntityCriteriaSupport;
import org.codefilarete.stalactite.engine.runtime.query.EntityQueryCriteriaSupport;
import org.codefilarete.stalactite.query.ConfiguredEntityCriteria;
import org.codefilarete.stalactite.query.EntityFinder;
import org.codefilarete.stalactite.query.builder.QuerySQLBuilderFactory.QuerySQLBuilder;
import org.codefilarete.stalactite.query.model.CriteriaChain;
import org.codefilarete.stalactite.query.model.GroupBy;
import org.codefilarete.stalactite.query.model.Having;
import org.codefilarete.stalactite.query.model.Limit;
import org.codefilarete.stalactite.query.model.OrderBy;
import org.codefilarete.stalactite.query.model.Query;
import org.codefilarete.stalactite.query.model.Select;
import org.codefilarete.stalactite.query.model.Selectable;
import org.codefilarete.stalactite.query.model.Where;
import org.codefilarete.stalactite.sql.ConnectionProvider;
import org.codefilarete.stalactite.sql.Dialect;
import org.codefilarete.stalactite.sql.ddl.structure.Column;
import org.codefilarete.stalactite.sql.ddl.structure.Table;
import org.codefilarete.stalactite.sql.result.Accumulator;
import org.codefilarete.stalactite.sql.result.ColumnedRow;
import org.codefilarete.stalactite.sql.result.ColumnedRowIterator;
import org.codefilarete.stalactite.sql.statement.PreparedSQL;
import org.codefilarete.stalactite.sql.statement.ReadOperation;
import org.codefilarete.stalactite.sql.statement.SQLExecutionException;
import org.codefilarete.stalactite.sql.statement.SQLOperation.SQLOperationListener;
import org.codefilarete.stalactite.sql.statement.SQLStatement;
import org.codefilarete.stalactite.sql.statement.StringParamedSQL;
import org.codefilarete.stalactite.sql.statement.binder.PreparedStatementWriter;
import org.codefilarete.stalactite.sql.statement.binder.ResultSetReader;
import org.codefilarete.tool.collection.Iterables;
import org.codefilarete.tool.collection.Maps;

import static org.codefilarete.tool.bean.Objects.preventNull;

/**
 * Class aimed at loading an entity graph which is selected by some criteria on some properties coming from a {@link CriteriaChain}.
 * 
 * Implementation is based on {@link EntityJoinTree} to build the query and the entity graph.
 * 
 * @author Guillaume Mary
 * @see EntityFinder#select(ConfiguredEntityCriteria, Map, OrderBy, Limit)
 */
public class RelationalEntityFinder<C, I, T extends Table<T>> implements EntityFinder<C, I> {
	
	private static final String PRIMARY_KEY_ALIAS = "rootId";
	
	private final EntityJoinTree<C, I> entityJoinTree;
	
	private final ConnectionProvider connectionProvider;

	private final Dialect dialect;
	private final EntityCriteriaSupport<C> criteriaSupport;
	private SelectExecutor<C, I> selectExecutor;
	
	private EntityTreeQuery<C> entityTreeQuery;
	
	private Query query;
	
	private SQLOperationListener<?> operationListener;
	
	public RelationalEntityFinder(EntityJoinTree<C, I> entityJoinTree,
								  SelectExecutor<C, I> selectExecutor,
								  ConnectionProvider connectionProvider,
								  Dialect dialect) {
		this.entityJoinTree = entityJoinTree;
		this.selectExecutor = selectExecutor;
		this.connectionProvider = connectionProvider;
		this.dialect = dialect;
		this.entityTreeQuery = new EntityTreeQueryBuilder<>(this.entityJoinTree, dialect.getColumnBinderRegistry()).buildSelectQuery();
		this.criteriaSupport = new EntityCriteriaSupport<>(this.entityJoinTree);
		
		PersisterBuilderContext.CURRENT.get().addBuildLifeCycleListener(new BuildLifeCycleListener() {
			@Override
			public void afterBuild() {
			}
			
			@Override
			public void afterAllBuild() {
				buildQuery();
			}
		});
	}
	
	public RelationalEntityFinder(AdvancedEntityPersister<C, I> mainPersister,
								  ConnectionProvider connectionProvider,
								  Dialect dialect,
								  boolean withImmediateQueryBuild) {
		this.entityJoinTree = mainPersister.getEntityJoinTree();
		this.connectionProvider = connectionProvider;
		this.dialect = dialect;
		this.entityTreeQuery = new EntityTreeQueryBuilder<>(this.entityJoinTree, dialect.getColumnBinderRegistry()).buildSelectQuery();
		this.criteriaSupport = new EntityCriteriaSupport<>(this.entityJoinTree, withImmediateQueryBuild);
	}
	
	private void buildQuery() {
		entityTreeQuery = new EntityTreeQueryBuilder<>(this.entityJoinTree, dialect.getColumnBinderRegistry()).buildSelectQuery();
		query = entityTreeQuery.getQuery();
	}
	
	@Override
	public void setOperationListener(SQLOperationListener<?> operationListener) {
		this.operationListener = operationListener;
	}
	
	@Override
	public EntityJoinTree<C, I> getEntityJoinTree() {
		return entityJoinTree;
	}
	
	@Override
	public EntityQueryCriteriaSupport<C, I> newCriteriaSupport() {
		return new EntityQueryCriteriaSupport<>(this, criteriaSupport.copy());
	}
	
	public Set<C> selectFromQueryBean(String sql, Map<String, Object> values) {
		// Computing parameter binders from values
		Map<String, PreparedStatementWriter<?>> parameterBinders = new HashMap<>();
		values.forEach((paramName, value) -> {
			PreparedStatementWriter<?> writer = dialect.getColumnBinderRegistry().getWriter(value.getClass());
			parameterBinders.put(paramName, writer);
		});
		
		return selectFromQueryBean(sql, values, parameterBinders);
	}
	
	public Set<C> selectFromQueryBean(String sql, Map<String, Object> values, Map<String, PreparedStatementWriter<?>> parameterBinders) {
		// we use EntityTreeQueryBuilder to get the inflater, please note that it also build the default Query
		EntityTreeInflater<C> inflater = entityTreeQuery.getInflater();
		
		// computing SQL readers from dialect binder registry
		Query query = entityTreeQuery.getQuery();
		Map<Selectable<?>, ResultSetReader<?>> selectParameterBinders = new HashMap<>();
		Map<Selectable<?>, String> aliases = new HashMap<>();
		query.getColumns().forEach(selectable -> {
			ResultSetReader<?> reader;
			String alias = preventNull(query.getAliases().get(selectable), selectable.getExpression());
			if (selectable instanceof Column) {
				reader = dialect.getColumnBinderRegistry().getReader((Column) selectable);
				selectParameterBinders.put(selectable, reader);
			} else {
				reader = dialect.getColumnBinderRegistry().getReader(selectable.getJavaType());
			}
			selectParameterBinders.put(selectable, reader);
			aliases.put(selectable, alias);
		});
		
		StringParamedSQL statement = new StringParamedSQL(sql, parameterBinders);
		statement.setValues(values);
		return new InternalExecutor(inflater, selectParameterBinders, aliases).execute(statement);
	}
	
	/**
	 * Implementation note: the load is done in 2 phases: one for root ids selection from criteria, a second for the whole graph load from found root ids.
	 */
	@Override
	public Set<C> select(ConfiguredEntityCriteria where,
						 Map<String, Object> valuesPerParam,
						 OrderBy orderBy,
						 Limit limit) {
		
		// When the condition contains some criteria on a collection, the ResultSet contains only data matching it,
		// then the graph is a partial view of the real entity. Therefore, when the condition contains some Collection criteria
		// we must load the graph in 2 phases: a first lookup for ids matching the result, and a second phase that loads the entity graph
		// according to the ids
		if (where.hasCollectionCriteria()) {
			Query queryClone = new Query(
					new Select(),
					entityTreeQuery.getQuery().getFromDelegate(),
					new Where<>(where.getCriteria()),
					new GroupBy(),
					new Having(),
					orderBy,
					limit);
			
			QuerySQLBuilder sqlQueryBuilder = dialect.getQuerySQLBuilderFactory().queryBuilder(queryClone);
			
			// First phase: selecting ids (made by clearing selected elements for performance issue)
			Column<T, I> pk = (Column<T, I>) Iterables.first(((Table) entityJoinTree.getRoot().getTable()).getPrimaryKey().getColumns());
			queryClone.select(pk, PRIMARY_KEY_ALIAS);
			Map<Column<?, ?>, String> aliases = Maps.asMap(pk, PRIMARY_KEY_ALIAS);
			Map<Column<?, ?>, ResultSetReader<?>> columnReaders = Maps.asMap(pk, dialect.getColumnBinderRegistry().getBinder(pk));
			Set<I> ids = readIds(sqlQueryBuilder.toPreparableSQL().toPreparedSQL(new HashMap<>()), columnReaders, aliases);
			
			if (ids.isEmpty()) {
				// No result found, we must stop here because request below doesn't support in(..) without values (SQL error from database)
				return Collections.emptySet();
			} else {
				// Second phase : selecting elements by found identifiers
				return selectExecutor.select(ids);
			}
		} else {
			// The condition doesn't have a criteria on a collection property (*-to-many): the load can be done with one query because the SQL criteria
			// doesn't make a subset of the entity graph
			// We clone the query to avoid polluting the instance one, else, from select(..) to select(..), we append the criteria at the end of it,
			// which makes the query usually returning no data (because of the condition mix)
			Query queryClone = new Query(
					new Select(entityTreeQuery.getQuery().getSelectDelegate()),
					entityTreeQuery.getQuery().getFromDelegate(),
					new Where<>(where.getCriteria()),
					new GroupBy(),
					new Having(),
					orderBy,
					limit);
			
			QuerySQLBuilder sqlQueryBuilder = dialect.getQuerySQLBuilderFactory().queryBuilder(queryClone);
			PreparedSQL preparedSQL = sqlQueryBuilder.toPreparableSQL().toPreparedSQL(valuesPerParam);
			return new InternalExecutor(entityTreeQuery).execute(preparedSQL);
		}
	}
	
	private Set<I> readIds(PreparedSQL preparedSQL, Map<Column<?, ?>, ResultSetReader<?>> columnReaders, Map<Column<?, ?>, String> aliases) {
		EntityInflater<C, I> entityInflater = entityJoinTree.getRoot().getEntityInflater();
		try (ReadOperation<Integer> closeableOperation = dialect.getReadOperationFactory().createInstance(preparedSQL, connectionProvider)) {
			ColumnedRowIterator rowIterator = new ColumnedRowIterator(closeableOperation.execute(), columnReaders, aliases);
			return Iterables.collect(() -> rowIterator, row -> entityInflater.giveIdentifier(row), HashSet::new);
		} catch (RuntimeException e) {
			throw new SQLExecutionException(preparedSQL.getSQL(), e);
		}
	}
	
	@Override
	public <R, O> R selectProjection(Consumer<Select> selectAdapter,
									 Map<String, Object> values,
									 Accumulator<? super Function<Selectable<O>, O>, Object, R> accumulator,
									 ConfiguredEntityCriteria where,
									 boolean distinct,
									 OrderBy orderBy,
									 Limit limit) {
		Query queryClone = new Query(new Select(), query.getFromDelegate(), new Where<>(where.getCriteria()), new GroupBy(), new Having(), orderBy, limit);
		queryClone.getSelectDelegate().setDistinct(distinct);
		QuerySQLBuilder sqlQueryBuilder = dialect.getQuerySQLBuilderFactory().queryBuilder(queryClone);
		
		// First phase : selecting ids (made by clearing selected elements for performance issue)
		selectAdapter.accept(queryClone.getSelectDelegate());
		Map<Selectable<?>, ResultSetReader<?>> columnReaders = Iterables.map(queryClone.getColumns(), Function.identity(), selectable -> dialect.getColumnBinderRegistry().getBinder(selectable.getJavaType()));
		
		PreparedSQL preparedSQL = sqlQueryBuilder.toPreparableSQL().toPreparedSQL(values);
		Map<Selectable<?>, String> aliases = queryClone.getAliases();
		
		// Propagating aliases from the original query to clone if the user didn't mention the alias
		// This avoids the later thrown exception "Column doesn't exist : null" because the column reading is based on the alias (which doesn't exist)
		Map<Selectable<?>, String> defaultAliases = query.getAliases();
		aliases.entrySet().forEach(alias -> {
			if (alias.getValue() == null) {
				alias.setValue(defaultAliases.get(alias.getKey()));
			}
		});
		return readProjection(preparedSQL, columnReaders, aliases, accumulator);
	}
	
	private <R, O> R readProjection(PreparedSQL preparedSQL,
									Map<Selectable<?>, ResultSetReader<?>> columnReaders,
									Map<Selectable<?>, String> aliases,
									Accumulator<? super Function<Selectable<O>, O>, Object, R> accumulator) {
		try (ReadOperation<Integer> closeableOperation = dialect.getReadOperationFactory().createInstance(preparedSQL, connectionProvider)) {
			ColumnedRowIterator rowIterator = new ColumnedRowIterator(closeableOperation.execute(), columnReaders, aliases);
			return accumulator.collect(Iterables.stream(rowIterator).map(row -> (Function<Selectable<O>, O>) row::get).collect(Collectors.toList()));
		} catch (RuntimeException e) {
			throw new SQLExecutionException(preparedSQL.getSQL(), e);
		}
	}
	
	/**
	 * Small class to avoid passing {@link EntityTreeQuery} as argument to all methods
	 */
	private class InternalExecutor {
		
		private final EntityTreeInflater<C> inflater;
		private final Map<Selectable<?>, ResultSetReader<?>> selectParameterBinders;
		private final Map<Selectable<?>, String> columnAliases;
		
		private InternalExecutor(EntityTreeQuery<C> entityTreeQuery) {
			this(entityTreeQuery.getInflater(), entityTreeQuery.getSelectParameterBinders(), entityTreeQuery.getColumnAliases());
		}
		
		public InternalExecutor(EntityTreeInflater<C> inflater,
								Map<Selectable<?>, ? extends ResultSetReader<?>> selectParameterBinders,
								Map<Selectable<?>, String> columnAliases) {
			this.inflater = inflater;
			this.selectParameterBinders = (Map<Selectable<?>, ResultSetReader<?>>) selectParameterBinders;
			this.columnAliases = columnAliases;
		}
		
		protected <ParamType> Set<C> execute(SQLStatement<ParamType> query) {
			try (ReadOperation<ParamType> readOperation = dialect.getReadOperationFactory().createInstance(query, connectionProvider)) {
				readOperation.setListener((SQLOperationListener<ParamType>) operationListener);
				// Note that setValues must be done after operationListener set
				readOperation.setValues(query.getValues());
				return transform(readOperation);
			} catch (RuntimeException e) {
				throw new SQLExecutionException(query.getSQL(), e);
			}
		}
		
		protected Set<C> transform(ReadOperation<?> closeableOperation) {
			ResultSet resultSet = closeableOperation.execute();
			// NB: we give the same ParametersBinders of those given at ColumnParameterizedSelect since the row iterator is expected to read column from it
			ColumnedRowIterator rowIterator = new ColumnedRowIterator(resultSet, selectParameterBinders, columnAliases);
			return transform(rowIterator);
		}
		
		protected Set<C> transform(Iterator<? extends ColumnedRow> rowIterator) {
			return inflater.transform(() -> (Iterator<ColumnedRow>) rowIterator, 50);
		}
	}
}