net.jrouter.paging.ibatis.PageInterceptor Maven / Gradle / Ivy

Go to download
package net.jrouter.paging.ibatis;

import java.util.*;
import java.util.concurrent.ConcurrentHashMap;
import java.util.function.BiFunction;
import java.util.stream.Collectors;
import java.util.stream.Stream;
import lombok.extern.slf4j.Slf4j;
import net.jrouter.paging.ibatis.delegate.SqlSourceDelegate;
import net.jrouter.paging.jdbc.dialect.Dialect;
import org.apache.ibatis.cache.CacheKey;
import org.apache.ibatis.executor.Executor;
import org.apache.ibatis.mapping.BoundSql;
import org.apache.ibatis.mapping.MappedStatement;
import org.apache.ibatis.mapping.ResultMap;
import org.apache.ibatis.plugin.*;
import org.apache.ibatis.session.ResultHandler;
import org.apache.ibatis.session.RowBounds;

/**
 * MyBatis的分页拦截器。需要传递{@code AbstractRowBounds}分页参数以甄别是否调用分页组件。
 *
 * @see org.apache.ibatis.session.Configuration#getMappedStatement(java.lang.String)
 * @see Executor#query(org.apache.ibatis.mapping.MappedStatement, java.lang.Object, org.apache.ibatis.session.RowBounds, org.apache.ibatis.session.ResultHandler)
 * @see AbstractRowBounds
 */
@Intercepts(
        {
                @Signature(type = Executor.class, method = "query",
                        args = {MappedStatement.class, Object.class, RowBounds.class, ResultHandler.class}),
                @Signature(type = Executor.class, method = "query",
                        args = {MappedStatement.class, Object.class, RowBounds.class, ResultHandler.class, CacheKey.class, BoundSql.class})
        }
)
@Slf4j
public class PageInterceptor implements Interceptor {

    /** 线程绑定的分页参数对象 */
    private static final ThreadLocal LOCAL_PAGE_ROW_BOUNDS = new ThreadLocal<>();

    /** {@code MappedStatement}对象缓存 */
    @lombok.Setter
    private Map mappedStatementCache = (Map) buildStatementCacheMap();

    /** 数据库语义 */
    @lombok.Setter
    private Dialect dialect = null;

    /**
     * 列表结果转换器。
     */
    @lombok.Setter
    private BiFunction pageListConverter;

    static {
        Set allFields = new HashSet<>(Arrays.asList(
                "resource",
                "configuration",
                "id",
                "fetchSize",
                "timeout",
                "statementType",
                "resultSetType",
                "sqlSource",
                "cache",
                "parameterMap",
                "resultMaps",
                "flushCacheRequired",
                "useCache",
                "resultOrdered",
                "sqlCommandType",
                "keyGenerator",
                "keyProperties",
                "keyColumns",
                "hasNestedResultMaps",
                "databaseId",
                "statementLog",
                "lang",
                "resultSets"
        ));
        //all fields check
        Stream.of(MappedStatement.class.getDeclaredFields())
                .map(field -> field.getName())
                .filter(s -> !allFields.contains(s))
                .forEach(s -> log.warn("Unknown field [{}] of class {}", s, MappedStatement.class));
    }

    //copy MappedStatement properties to another
    private static void cloneMappedStatement(MappedStatement.Builder builder, MappedStatement ms) {
        builder.cache(ms.getCache());
        builder.databaseId(ms.getDatabaseId());
        builder.fetchSize(ms.getFetchSize());
        builder.flushCacheRequired(ms.isFlushCacheRequired());
        builder.keyGenerator(ms.getKeyGenerator());
        builder.keyProperty(arrayToDelimitedString(ms.getKeyProperties()));
        builder.lang(ms.getLang());
        builder.parameterMap(ms.getParameterMap());
        builder.resource(ms.getResource());
        builder.resultMaps(ms.getResultMaps());
        builder.resultOrdered(ms.isResultOrdered());
        builder.resultSets(arrayToDelimitedString(ms.getResultSets()));
        builder.resultSetType(ms.getResultSetType());
        builder.statementType(ms.getStatementType());
        builder.timeout(ms.getTimeout());
        builder.useCache(ms.isUseCache());
    }

    /**
     * @see MappedStatement#delimitedStringToArray(String)
     */
    private static String arrayToDelimitedString(String[] array) {
        if (array == null || array.length == 0) {
            return null;
        }
        return Stream.of(array).collect(Collectors.joining(","));
    }

    @Override
    public Object intercept(Invocation invocation) throws Throwable {
        final Object[] args = invocation.getArgs();
        AbstractRowBounds pageRowBounds = null;
        if (args[2] instanceof AbstractRowBounds && dialect != null) {
            MappedStatement ms = (MappedStatement) args[0];
            Object parameter = args[1];
            pageRowBounds = (AbstractRowBounds) args[2];
            //ResultHandler resultHandler = (ResultHandler) args[3];
            LOCAL_PAGE_ROW_BOUNDS.set(pageRowBounds);

            MappedStatementCache cache = mappedStatementCache.get(ms.getId());
            if (cache == null) {
                cache = new MappedStatementCache();
                MappedStatement.Builder pageBuilder = new MappedStatement.Builder(
                        ms.getConfiguration(),
                        ms.getId() + "##Page",
                        new SqlSourceDelegate.Page(ms.getConfiguration(), ms.getSqlSource(), dialect) {

                            @Override
                            public AbstractRowBounds getPageRowBounds() {
                                return LOCAL_PAGE_ROW_BOUNDS.get();
                            }
                        }, ms.getSqlCommandType());
                cloneMappedStatement(pageBuilder, ms);
                MappedStatement.Builder countBuilder = new MappedStatement.Builder(
                        ms.getConfiguration(),
                        ms.getId() + "##Count",
                        new SqlSourceDelegate.Count(ms.getConfiguration(), ms.getSqlSource(), dialect) {

                            @Override
                            public AbstractRowBounds getPageRowBounds() {
                                return LOCAL_PAGE_ROW_BOUNDS.get();
                            }
                        }, ms.getSqlCommandType());
                cloneMappedStatement(countBuilder, ms);
                //count查询
                List resultMaps = new ArrayList<>();
                ResultMap resultMap = new ResultMap.Builder(ms.getConfiguration(), ms.getId() + "##Count", int.class, Collections.EMPTY_LIST).build();
                resultMaps.add(resultMap);
                countBuilder.resultMaps(resultMaps);

                cache.pageMs = pageBuilder.build();
                cache.countMs = countBuilder.build();
                mappedStatementCache.put(ms.getId(), cache);
            }
            args[0] = cache.pageMs;
            args[2] = RowBounds.DEFAULT;
            if (pageRowBounds.isAutoCount()) {
                Executor executor = (Executor) invocation.getTarget();
                List countRes = executor.query(cache.countMs, parameter, RowBounds.DEFAULT, null);
                int totalCount = countRes.get(0).intValue();
                pageRowBounds.countTotalElements(totalCount);
                if (totalCount == 0) {
                    pageRowBounds.setContent(Collections.EMPTY_LIST);
                    LOCAL_PAGE_ROW_BOUNDS.remove();
                    return Collections.EMPTY_LIST;
                }
            }
        }
        try {
            Object res = invocation.proceed();
            if (pageRowBounds != null && (res instanceof List)) {
                pageRowBounds.setContent((List) res);
                if (pageListConverter != null) {
                    return pageListConverter.apply(pageRowBounds, (List) res);
                }
            }
            return res;
        } finally {
            LOCAL_PAGE_ROW_BOUNDS.remove();
        }
    }

    /**
     * Use map as statement cache.
     *
     * @return {@code ConcurrentHashMap} as cache.
     */
    protected Map buildStatementCacheMap() {
        return new ConcurrentHashMap<>();
    }

    @Override
    public Object plugin(Object target) {
        if (target instanceof Executor) {
            return Plugin.wrap(target, this);
        } else {
            return target;
        }
    }

    @Override
    public void setProperties(Properties properties) {
    }

    //MappedStatement缓存对象
    private static final class MappedStatementCache {

        MappedStatement pageMs;

        MappedStatement countMs;

    }
}