org.tkit.quarkus.jpa.daos.PagedQuery Maven / Gradle / Ivy

Go to download
package org.tkit.quarkus.jpa.daos;

import java.util.stream.Stream;

import jakarta.persistence.EntityManager;
import jakarta.persistence.NoResultException;
import jakarta.persistence.criteria.*;

import org.hibernate.query.sqm.tree.SqmCopyContext;
import org.hibernate.query.sqm.tree.select.SqmQueryPart;
import org.hibernate.query.sqm.tree.select.SqmSelectStatement;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.tkit.quarkus.jpa.exceptions.DAOException;

/**
 * The page query.
 *
 * @param  the entity class.
 */
public class PagedQuery {
    /**
     * The logger for this class.
     */
    private static final Logger log = LoggerFactory.getLogger(PagedQuery.class);
    /**
     * The entity manager.
     */
    private EntityManager em;

    /**
     * The search criteria.
     */
    private CriteriaQuery criteria;

    /**
     * The search count criteria.
     */
    private CriteriaQuery countCriteria;

    /**
     * The current page.
     */
    private Page page;

    /**
     * Default constructor.
     *
     * @param em the entity manager.
     * @param criteria the search criteria
     * @param page the start page.
     */
    public PagedQuery(EntityManager em, CriteriaQuery criteria, Page page, String idAttributeName) {
        this.em = em;
        this.criteria = setDefaultSorting(em, criteria, idAttributeName);
        this.page = page;
        this.countCriteria = createCountCriteria(em, criteria);
    }

    public PageResult getPageResult() {
        try {

            // get count
            var q = em.createQuery(countCriteria);
            Long count = q.getSingleResult();
            // return empty page for count zero
            if (count == 0) {
                return PageResult.empty();
            }

            // get stream
            var sq = em.createQuery(criteria);
            Stream stream = sq
                    .setFirstResult(page.number() * page.size())
                    .setMaxResults(page.size())
                    .getResultStream();
            // create page result
            return new PageResult<>(count, stream, page);
        } catch (NoResultException nex) {
            // return empty page for getSingleResult() throws NoResultException
            return PageResult.empty();
        } catch (Exception ex) {
            String entityClass = criteria.getResultType() != null ? criteria.getResultType().getName() : null;
            throw new DAOException(Errors.GET_PAGE_RESULT_ERROR, ex, page.number(), page.size(), entityClass);
        }
    }

    private static  CriteriaQuery setDefaultSorting(EntityManager em, CriteriaQuery criteria, String idAttributeName) {
        Root root = null;
        try {
            CriteriaBuilder builder = em.getCriteriaBuilder();
            if (criteria.getOrderList().isEmpty()) {
                log.warn(
                        "Paged query used without explicit orderBy. Ordering of results between pages not guaranteed. Please add an orderBy clause to your query.");
                root = findRoot(criteria, criteria.getResultType());
                if (root != null) {
                    criteria.orderBy(builder.asc(root.get(idAttributeName)));
                    log.warn("Default sorting by '{}' attribute is added.", idAttributeName);
                }
            }
        } catch (IllegalArgumentException ex) {
            log.error("There is no id attribute for Root:{}", root);
        }
        return criteria;
    }

    /**
     * Gets the current page.
     *
     * @return the current page.
     */
    public Page getPage() {
        return page;
    }

    /**
     * Gets the search count criteria.
     *
     * @return the search count criteria.
     */
    public CriteriaQuery countCriteria() {
        return countCriteria;
    }

    /**
     * Gets the search criteria.
     *
     * @return the search criteria.
     */
    public CriteriaQuery criteria() {
        return criteria;
    }

    /**
     * Move to the previous page.
     *
     * @return the page query.
     */
    public PagedQuery previous() {
        if (page.number() > 0) {
            page = Page.of(page.number() - 1, page.size());
        }
        return this;
    }

    /**
     * Move to the next page.
     *
     * @return the page query.
     */
    public PagedQuery next() {
        page = Page.of(page.number() + 1, page.size());
        return this;
    }

    /**
     * Internal error code.
     */
    public enum Errors {

        /**
         * Gets the page result error.
         */
        GET_PAGE_RESULT_ERROR;
    }

    /**
     * Create a row count CriteriaQuery from a CriteriaQuery
     *
     * @param em entity manager
     * @param criteria source criteria
     * @param  the entity type.
     * @return row count CriteriaQuery
     */
    public static  CriteriaQuery createCountCriteria(EntityManager em, CriteriaQuery criteria) {
        CriteriaBuilder builder = em.getCriteriaBuilder();
        CriteriaQuery countCriteria = createCountCriteriaQuery(builder, criteria);
        Expression countExpression;
        Root root = findRoot(countCriteria, criteria.getResultType());
        if (criteria.isDistinct()) {
            countExpression = builder.countDistinct(root);
        } else {
            countExpression = builder.count(root);
        }
        return countCriteria.select(countExpression);
    }

    /**
     * Creates count criteria query base on the {@code from} criteria query..
     *
     * @param builder the criteria builder.
     * @param from source Criteria.
     * @return count criteria query.
     */
    @SuppressWarnings("unchecked")
    public static  CriteriaQuery createCountCriteriaQuery(CriteriaBuilder builder, CriteriaQuery from) {
        CriteriaQuery result = builder.createQuery(Long.class);
        SqmSelectStatement copy = ((SqmSelectStatement) from).copy(SqmCopyContext.simpleContext());
        SqmSelectStatement r = (SqmSelectStatement) result;
        SqmQueryPart part = copy.getQueryPart();
        // remove group by from the count
        part.setOrderByClause(null);
        r.setQueryPart(part);
        r.getOrderList().clear();
        return result;
    }

    /**
     * Find the Root with type class on {@link CriteriaQuery} Root Set for the {@code clazz}.
     *
     * @param query criteria query
     * @param clazz root type
     * @param  the type of the root class.
     * @return the root of the criteria query or {@code null} if none
     */
    @SuppressWarnings("unchecked")
    public static  Root findRoot(CriteriaQuery query, Class clazz) {
        for (Root r : query.getRoots()) {
            if (clazz.equals(r.getJavaType())) {
                return (Root) r;
            }
        }
        return null;
    }

    /**
     * Copy Joins
     *
     * @param from source Join
     * @param to destination Join
     */
    public static void copyJoins(From from, From to, AliasCounter counter) {
        from.getJoins().forEach(join -> {
            Join item = to.join(join.getAttribute().getName(), join.getJoinType());
            item.alias(createAlias(join, counter));
            copyJoins(join, item, counter);
        });
    }

    /**
     * Copy Fetches
     *
     * @param from source From
     * @param to destination From
     */
    public static void copyFetches(From from, From to) {
        from.getFetches().forEach(fetch -> {
            Fetch item = to.fetch(fetch.getAttribute().getName(), fetch.getJoinType());
            copyFetches(fetch, item);
        });
    }

    /**
     * Copy Fetches
     *
     * @param from source From
     * @param to destination From
     */
    public static void copyFetches(Fetch from, Fetch to) {
        from.getFetches().forEach(fetch -> {
            Fetch item = to.fetch(fetch.getAttribute().getName(), fetch.getJoinType());
            copyFetches(fetch, item);
        });
    }

    /**
     * Gets The result alias, if none set a default one and return it
     *
     * @param selection the selection
     * @return root alias or generated one
     */
    public static  String createAlias(Selection selection, AliasCounter counter) {
        String alias = selection.getAlias();
        if (alias == null) {
            alias = counter.next();
            selection.alias(alias);
        }
        return alias;

    }

    /**
     * Criteria query alias copy counter.
     */
    public static class AliasCounter {

        private long index = 0;

        public String next() {
            return "a_" + index++;
        }
    }

    @Override
    public String toString() {
        return "PagedQuery{" +
                "page=" + page +
                '}';
    }
}