Implementing JPAQuery.count(EntityPath<?>) for QueryDSL

BetterJPAQuery.java

/*
 * Copyright 2011, Mysema Ltd
 * 
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 * http://www.apache.org/licenses/LICENSE-2.0
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

package com.anjlab.querydsl;

import java.util.regex.MatchResult;
import java.util.regex.Matcher;
import java.util.regex.Pattern;

import javax.persistence.EntityManager;

import com.mysema.query.DefaultQueryMetadata;
import com.mysema.query.QueryMetadata;
import com.mysema.query.jpa.JPQLTemplates;
import com.mysema.query.jpa.impl.AbstractJPAQuery;
import com.mysema.query.jpa.impl.JPAProvider;
import com.mysema.query.types.EntityPath;

/**
 * JPAQuery is the default implementation of the JPQLQuery interface for JPA
 *
 * @author tiwe
 * @author dmitrygusev
 *
 */
public final class BetterJPAQuery extends AbstractJPAQuery<BetterJPAQuery> {

    /**
     * Creates a new detached query
     * The query can be attached via the clone method
     */
    public BetterJPAQuery() {
        super(null, JPQLTemplates.DEFAULT, new DefaultQueryMetadata());
    }

    /**
     * Creates a new EntityManager bound query
     *
     * @param em
     */
    public BetterJPAQuery(EntityManager em) {
        super(em, JPAProvider.getTemplates(em), new DefaultQueryMetadata());
    }
    
    /**
     * Creates a new EntityManager bound query
     *
     * @param em
     */
    public BetterJPAQuery(EntityManager em, QueryMetadata metadata) {
        super(em, JPAProvider.getTemplates(em), metadata);
    }

    /**
     * Creates a new query
     *
     * @param em
     * @param patterns
     */
    public BetterJPAQuery(EntityManager em, JPQLTemplates patterns) {
        super(em, patterns, new DefaultQueryMetadata());
    }

    /**
     * Creates a new query
     * 
     * @param em
     * @param templates
     * @param metadata
     */
    public BetterJPAQuery(EntityManager em, JPQLTemplates templates, QueryMetadata metadata) {
        super(em, templates, metadata);
    }

    /**
     * Clone the state of this query to a new JPAQuery instance with the given EntityManager
     *
     * @param entityManager
     * @return
     */
    public BetterJPAQuery clone(EntityManager entityManager) {
        BetterJPAQuery q = new BetterJPAQuery(entityManager, JPAProvider.getTemplates(entityManager), getMetadata().clone());
        q.projection = projection;
        q.flushMode = flushMode;
        q.hints.putAll(hints);
        q.lockMode = lockMode;
        return q;
    }
    
    /*
     * Overriding select count(...)
     */
    
    private static final Pattern SELECT_COUNT_PATTERN = Pattern.compile("select count\\((?:distinct )?([^\\)]+)\\)");
    
    @Override
    protected String buildQueryString(boolean forCountRow)
    {
        String queryString = super.buildQueryString(forCountRow);
        
        if (forCountRow && entityPath != null)
        {
            Matcher matcher = SELECT_COUNT_PATTERN.matcher(queryString);
            
            if (matcher.find())
            {
                MatchResult matchResult = matcher.toMatchResult();
                
                String target = entityPath.getMetadata().getName();
                
                StringBuilder builder = new StringBuilder(queryString.length() + target.length());
                
                queryString =  builder.append(queryString)
                                      .replace(matchResult.start(1), matchResult.end(1), target).toString();
            }
        }
        
        return queryString;
    }
    
    private EntityPath<?> entityPath;
    
    public long count(EntityPath<?> entityPath)
    {
        this.entityPath = entityPath;
        
        return super.count();
    }
}