# 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.api.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.api.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().getFrom(),
    					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().getSelect()),
    					entityTreeQuery.getQuery().getFrom(),
    					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.getFrom(), new Where(where.getCriteria()), new GroupBy(), new Having(), orderBy, limit);
    		queryClone.getSelect().setDistinct(distinct);
    		QuerySQLBuilder sqlQueryBuilder = dialect.getQuerySQLBuilderFactory().queryBuilder(queryClone);

    		// First phase : selecting ids (made by clearing selected elements for performance issue)
    		selectAdapter.accept(queryClone.getSelect());
    		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);
    		}
    	}
    }